Learning Task-Specific Antibody Representations via Function-Aware Masking

# Learning Task-Specific Antibody Representations via Function-Aware Masking
Source: [https://arxiv.org/html/2609.00518](https://arxiv.org/html/2609.00518)
Thomas A\. WaltonAmirali AghazadehAffiliation:School of Electrical and Computer Engineering, Georgia Institute of TechnologyCorresponding Author:amiralia@gatech\.edu\*Equal contributionAffiliation:School of Computer Science, Georgia Institute of Technology

###### Abstract

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\.9×\\timesimprovement 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\.

## 1Introduction

The development of therapeutic antibodies sits at the forefront of drug discovery, providing targeted treatments for complex oncological, autoimmune, and infectious diseases\. Learning expressive sequence representations is central to optimizing these biologics for clinical viability\. Antibody\-specific language models trained under the masked language modeling \(MLM\) objective \(abLMs\), such as AntiBERTa\[[1](https://arxiv.org/html/2609.00518#bib.bib2),[2](https://arxiv.org/html/2609.00518#bib.bib3)\]and AbLang\[[3](https://arxiv.org/html/2609.00518#bib.bib4),[4](https://arxiv.org/html/2609.00518#bib.bib5)\], encode representations that demonstrate strong generalization across diverse sequence design and optimization tasks\.

Improvements to these models have largely stemmed from scaling pretraining data and refining architecture\. However, recent work has demonstrated that the corruption process itself serves as an inductive bias that shapes how learned representations generalize across tasks\[[5](https://arxiv.org/html/2609.00518#bib.bib1),[6](https://arxiv.org/html/2609.00518#bib.bib16),[7](https://arxiv.org/html/2609.00518#bib.bib17)\]\. Specifically, preferential masking of antibody complementarity\-determining regions \(CDRs\) increases transfer performance on CDR infilling, binding affinity, and binding specificity tasks\.

Masking of CDRs, however, targets only one facet of antibody function\. Therapeutic antibodies must also fold stably, resist aggregation, and exhibit mutational robustness to ensure clinical developability\. The residues governing these properties often lie outside of CDRs; they are distributed across the framework region that provides the structural scaffold\. The success of CDR masking and its variants points to a broader principle and raises a more general question: can antibody masking strategies be designed using any arbitrary functional prior to align representations with their corresponding downstream tasks?

![Refer to caption](https://arxiv.org/html/2609.00518v1/fig1.png)Figure 1:Function\-aware masking concentrates representation learning on critical antibody regions\.a,The rigid structure of an antibody\-antigen complex \(PDB ID: 3HFM\), illustrating that functional variation is heavily localized to a sparse subset of sequence locations\.b,Masked language models \(MLMs\) trained on hundreds of thousands of antibodies implicitly learn how sequence variation affects function through a corruption \(masking\) and reconstruction objective\.c,Random masking treats all sequence positions as relevant; however, sequence positions most consequential to downstream fitness are localized to few regions\. Function\-aware masking leverages functional annotations to guide mask placement, focusing the training signal on these relevant regions\.To address this question, we introduce function\-aware masking \(Figure[1](https://arxiv.org/html/2609.00518#S1.F1)\), a family of masking strategies that bias mask placement toward the residues most associated with a given functional task\. Each strategy scores residues using a functional signal derived from either sequence labels or model predictions \(e\.g\., IMGT labels or IgFold predictions\[[8](https://arxiv.org/html/2609.00518#bib.bib12)\]\) and preferentially masks those positions during training \(Figure[1](https://arxiv.org/html/2609.00518#S1.F1)c\)\. Importantly, function\-aware masking shapes only the training objective: resulting models remain sequence\-only at inference and incur no additional parameter cost\. These masking strategies, hereinafter referred to as specialist strategies, improve performance over random masking on tasks that align with their prior\. To consistently drive improvements across multiple functional axes simultaneously, we further introduce hybrid masking strategies that combine functional priors through stochastic mixtures and learning curricula\. We evaluate function\-aware masking in the pretraining setting and assess performance over five downstream tasks spanning binding, structure, developability, and generative objectives\. Our contributions are as follows:

1. 1\.We introduce function\-aware antibody masking, a class of masking algorithms that impose functional inductive biases guided by the downstream task of interest\. We open\-source our implementation here:[https://github\.com/amirgroup\-codes/function\-aware\-masking\.git](https://github.com/amirgroup-codes/function-aware-masking.git)
2. 2\.We demonstrate that specialized masking strategies, focusing on an individual biological prior, improve performance over random masking on tasks aligning with their prior\. Specialist strategies improved performance on structure\-related tasks by up to14%14\\%and CDR\-related tasks by up to5\.9×5\.9\\times\.
3. 3\.To further improve generalization, we demonstrate that hybrid masking strategies increase representation quality across multiple functional axes, mitigating performance degradation exhibited by specialists on unaligned tasks\.

## 2Background

Corruption is an essential component of MLM training\. The most common corruption scheme, replacing tokens with a mask token \(\[MASK\]\), has historically been applied uniformly at random across sequence positions with a fixed probability, typically 15%\[[9](https://arxiv.org/html/2609.00518#bib.bib6)\]\. Foundational work in natural language processing demonstrated that mask placement can influence learned representations, priming them for tasks related to inferring semantic structure\[[10](https://arxiv.org/html/2609.00518#bib.bib18),[11](https://arxiv.org/html/2609.00518#bib.bib19),[12](https://arxiv.org/html/2609.00518#bib.bib20)\]\. We recently extended this idea to protein sequences by guiding mask placement based on structural contacts, finding that resulting representations exhibited superior generalization on a range of extrapolation tasks\[[7](https://arxiv.org/html/2609.00518#bib.bib17)\]\.

Biological priors in antibodies\.While general proteins benefit from structural priors, antibodies exhibit a highly specialized architecture governed by multiple distinct functional drivers\. Antigen recognition is primarily mediated by six CDRs, forming contiguous, highly variable spans\. Within these loops, a specific subset of residues \(paratopes\) forms a direct interface with the target\. During affinity maturation, these binding interfaces and their corresponding framework regions undergo somatic hypermutation \(SHM\), introducing variations relative to the inherited germline sequence\. Beyond binding, the overall stable fold of the antibody is maintained by a network of structural contacts, including long\-range interactions that form epistatic hubs and dictate conformation\. Notably, the aforementioned functional drivers constitute a small proportion of antibody sequences\. It follows that much of the corruption applied during MLM training outside of these functionally critical regions is trivially easy to decode given the surrounding context\.

Antibody masking\.Previous works have explored masking strategies for antibodies, primarily focusing on CDRs\. Ng & Briney\[[5](https://arxiv.org/html/2609.00518#bib.bib1)\]introduced preferential masking of non\-templated CDR3, finding that learned representations generalized better to CDR3 residue recovery and improved binding specificity predictions\. Talaei et al\.\[[6](https://arxiv.org/html/2609.00518#bib.bib16)\]extended this work by introducing a hybrid CDR and framework masking strategy, further demonstrating improvements on binding affinity predictions\. Whereas previous works focused predominantly on CDRs, this work explores masking strategies for a larger set of function\-related tasks aligned with the previously discussed biological priors\.

## 3Methods

We introduce function\-aware masking, a set of masking algorithms leveraging biological priors of antibodies to enable representations across an array of tasks\. Function\-aware masking is split into two categories: specialist strategies focused primarily on one task, and hybrid strategies that mix priors from many tasks\.

Mask formulation and precomputed priors\.To enable targeted mask sampling while limiting computational overhead, we precompute per\-residue biological labels and cache them alongside the training data\. Labels include CDR locations, paratope probabilities, germline mutation status, and structural topology\. During data collation, a mask setMMis sampled for sequenceSS\(excluding \[CLS\]/\[SEP\] tokens\)\. To isolate the effect of any given strategy compared to the rest, we fix the mask rate at0\.150\.15\. The selected positions then undergo the standard MLM corruption procedure introduced by BERT\[[9](https://arxiv.org/html/2609.00518#bib.bib6)\]:80%80\\%of tokens are replaced by a \[MASK\] token,10%10\\%by a random amino acid, and the remaining10%10\\%remain unchanged\.

### 3\.1Specialist Strategies

The goal of specialist masking strategies is to align representations learned by an antibody MLM with a specific downstream task or functional property\. For the weight\-based strategies, the mask sampler assigns each position a weightwiw\_\{i\}encoding its relevance to the target property\. Weights are then normalized into per\-position masking probabilities calibrated to an expected budget of15%15\\%\(Equation[1](https://arxiv.org/html/2609.00518#S3.E1)\)\. Mask locations are determined with an independent Bernoulli draw \(Equation[2](https://arxiv.org/html/2609.00518#S3.E2)\)\. For non\-weight\-based strategies which depend on previously placed masks \(span,structure,structure\-lr\), masks are placed iteratively until a fixed budget of 15% is reached\.

pi=min\(0\.15wiw¯,1\)1\[i∈S\],w¯=1\|S\|∑j∈Swj,𝔼\[\|M\|\]=0\.15\|S\|\.p\_\{i\}\\;=\\;\\min\\\!\\Bigl\(0\.15\\,\\tfrac\{w\_\{i\}\}\{\\bar\{w\}\},\\,1\\Bigr\)\\,\\mathbf\{1\}\[i\\in S\],\\qquad\\bar\{w\}=\\tfrac\{1\}\{\|S\|\}\\\!\\sum\_\{j\\in S\}w\_\{j\},\\qquad\\mathbb\{E\}\[\|M\|\]=0\.15\\,\|S\|\.\(1\)mi∼Bernoulli⁡\(pi\)​independently,M=\{i:mi=1\}\.m\_\{i\}\\sim\\mathrm\{Bernoulli\}\(p\_\{i\}\)\\ \\text\{ independently\},\\qquad M=\\\{\\,i:m\_\{i\}=1\\,\\\}\.\(2\)
Under this formulation, we evaluate seven specialist strategies designed to target distinct functional, evolutionary, and structural priors\. Details on exact implementations can be found in Section[4](https://arxiv.org/html/2609.00518#S4)\.

cdr: Targets hypervariable regions by over\-sampling the complementarity\-determining loops \(wFW:wCDR1:wCDR2:wCDR3=:::6w\_\{\\text\{FW\}\}\\\!:\\\!w\_\{\\text\{CDR1\}\}\\\!:\\\!w\_\{\\text\{CDR2\}\}\\\!:\\\!w\_\{\\text\{CDR3\}\}\\\!=\\\!1\\\!:\\\!3\\\!:\\\!3\\\!:\\\!6\), generalizing CDR3 masking to the remaining CDRs\[[5](https://arxiv.org/html/2609.00518#bib.bib1)\]\.

span: Topology\-agnostic, masks contiguous spans of sequence positions using a geometric distribution \(p=0\.2p\\\!=\\\!0\.2,ℓmax=10\\ell\_\{\\text\{max\}\}\\\!=\\\!10\)\.

interface: Targets the physical binding interface by placing a6:16:1weight ratio on paratope\-contacting residues\.

germline: Targets the evolutionary trajectory of the antibody by placing a6:16:1weight ratio on germline\-mutated residues\.

intersection: A joint specialist that multiplies the paratope and germline weights:wi=wipara⋅wigermw\_\{i\}=w^\{\\text\{para\}\}\_\{i\}\\cdot w^\{\\text\{germ\}\}\_\{i\}\. This strategy targets residues responsible for affinity maturation, concentrating masks on positions that are both antigen\-contacting and somatically mutated from the germline\.

structure: A 3D\-aware scheme that weights residues based on predicted spatial proximity\. Utilizes the structure’s five nearest neighbors to mask structural neighborhoods\. Masks are sampled across neighborhoods, leaving enough spatial context to decode the masked residues\.

structure\-lr: An extension ofstructurethat applies a stricter partitioning rule: only structural couplings greater than four sequence indices away are considered\. Prioritizes modeling long\-range contacts during reconstruction\.

### 3\.2Hybrid Strategies

Specialist strategies target one desired antibody property; however, there may be many such properties to optimize\. To resolve this, we develop hybrid strategies that stochastically mix functional priors over the course of training\. Given a set ofKKmasking strategies, we define the mixing distributionπ⁡\(t\)=\(p1​\(t\),…,pK​\(t\)\)\\pi\(t\)=\(p\_\{1\}\(t\),\\dots,p\_\{K\}\(t\)\), wherepi​\(t\)p\_\{i\}\(t\)represents the probability of drawing specialist strategyiiat training steptt\(∑sps​\(t\)=1\\sum\_\{s\}p\_\{s\}\(t\)=1\)\. Mixing distributions follow two principles: masks should be sampled uniformly before specializing\[[13](https://arxiv.org/html/2609.00518#bib.bib26)\], and the best performing specialists should have the highest probability of being sampled\. Sampling is performed by drawing strategyi∼π⁡\(t\)i\\sim\\pi\(t\)and applying masks per sequence\. We dynamically adjustπ⁡\(t\)\\pi\(t\)over the course of training to bias which functional priors are drawn\. Details on mixing distribution construction are available in Appendix[B\.2](https://arxiv.org/html/2609.00518#A2.SS2)\.

Hybrid strategies interpolate between the following specialists unless specified otherwise, referenced in the following order: \[random,cdr,span,structure\-lr,interface,germline\]\. We evaluate the following hybrid strategies:

hybrid:Basis for the following hybrid strategies\. Emphasizesrandomandspanmasking before gradually shifting towards specialists\.

hybrid\-random:A randomized hybrid control strategy which samples fromπ⁡\(t\)\\pi\(t\)uniformly\.

hybrid\-stretched:Skews the mixing distribution heavily towardrandomandspanto start\. Ends with the same distribution ashybrid\. Delays transitions in mixing across more time steps\.

hybrid\-reverse:Reverses the direction in whichhybrid\-stretchedis applied\.

hybrid\-weighted:Starts with a general strategy similar tohybrid\-stretched, shifts focus tointerfacemasking, followed by a mix ofgermlineandinterfacemasking, before finally ending with a more balanced distribution\.

hybrid\-perbatch:Identical schedule tohybrid, but draws a single specialist per batch of sequences instead of per sequence\.

hybrid\-warmstart:Continued pretraining instead of a curriculum\. Takes a fully trainedinterfacemasking model and trains for an additional 50,000 steps on a balanced mixing distribution\.

## 4Experiments

We study function\-aware masking in the pretraining setting, training models parameterized by the RoFormer architecture\[[14](https://arxiv.org/html/2609.00518#bib.bib15)\]\. Specifically, we instantiate a medium\-sized RoFormer model \(38\.1M parameters\) utilizing the AntiBERTa2 tokenizer\[[2](https://arxiv.org/html/2609.00518#bib.bib3)\], AdamW optimizer \(peak LR×10−55\\\!\\times\\\!10^\{\-5\}, 5 % warmup, cosine decay\), batch size 64, fp16, and 125,000 training steps\. All masking strategies are evaluated using this setup with the exception ofhybrid\-warmstart, which trains for an additional 50,000 steps on top of theinterfacemodel checkpoint at a reduced learning rate \(×10−52\\\!\\times\\\!10^\{\-5\}\)\. Each model is trained three times across different initialization seeds\. Pretraining runs take approximately five hours on one RTX A6000 under fp16;hybrid\-warmstartadds 2 hours per seed\. One time IgFold structure prediction required≈\\approx240 GPU hours\.

Data\.We sampled 500,000 heavy\-chain variable\-domain \(VH\) sequences from the Observed Antibody Space\[[15](https://arxiv.org/html/2609.00518#bib.bib10)\]and filtered them to the 20 canonical amino acids and lengths in\[80,160\]\[80,160\], leaving 497,309 total sequences\. Models are pretrained on the same 90/10 training and hold\-out split across each run\. CDR1/2/3 boundaries are taken from the OAS IMGT annotations, with ANARCI\[[16](https://arxiv.org/html/2609.00518#bib.bib8)\]as a fallback for unannotated sequences\. Per\-residue paratope probabilities are predicted by a teacher model: AntiBERTa2 fine\-tuned with a per\-token classification head\. The classification head predicts over the TDC SAbDab\_Liberis paratope set\[[17](https://arxiv.org/html/2609.00518#bib.bib13),[18](https://arxiv.org/html/2609.00518#bib.bib9)\], storing residue labels in\[0,1\]\[0,1\]\. Germline labels are computed against a per\-gene consensus built from the training corpus; sequences are grouped by V\- and J\-gene calls, with position\-wise majority voting forming a consensus for genes represented by at least 20 sequences\. A residue is labeled mutated \(1\.01\.0\) if it differs from the consensus, germline \(0\.00\.0\) if it matches, or CDR3 junction \(0\.50\.5\) where no consensus is available\. Structures are predicted with IgFold\[[8](https://arxiv.org/html/2609.00518#bib.bib12)\]\(experimental X\-ray structures exist for<1%<1\\,\\%of OAS\), requiring approximately 240 GPU\-hours\. We build a per\-residuekk\-nearest neighbors graph usingCα−CαC\_\{\\alpha\}\-C\_\{\\alpha\}distances; structure\-aware strategies mask residues within these neighborhoods \(k=5k=5\)\.structure\-lrmasking follows this logic but restricts neighborhoods to residues more than four indices apart in sequence, targeting long\-range interactions\. Datasets used for paratope and structure prediction are presplit to ensure no leakage with the evaluation sequences\.

Evaluation\.We evaluate learned representations in two stages: first, the model is pretrained with a given masking strategy; then, the pretrained encoder is frozen and embeddings are extracted from the last layer as input to downstream probes\. We assess seven metrics designed to evaluate the diverse biological functions of antibodies, with each metric isolating a distinct functional property captured by learned representations\.

To evaluate generative recovery of the antigen\-binding loop, we measure exact match CDR3 infilling \(CDR3\) zero\-shot on held\-out OAS heavy chains\. All remaining metrics are evaluated via linear or bilinear probing\. For functional binding, paratope AUPRC \(Para\) and MCC evaluate the identification of antigen\-contacting residues on TDC SAbDab\_Liberis \(4\.5 Å contact threshold\)\. Paratope AUPRC is evaluated using a per\-token linear classifier trained with class\-weighted binary cross\-entropy; the MCC decision threshold is fit on the validation set by maximizing Youden’s J\. To assess structural functions, contact map AUROC and long\-range precision\-at\-L \(Cont; sequence separation≥\\geq24\) measure tertiary fold recovery on SAbDab crystal structures \(8 ÅCαC\_\{\\alpha\}\) utilizing a bilinear head over residue\-pair embeddings\. The structure probe Spearmanρ\\rho\(Struct\) assesses how linearly the embedding space encodes three\-dimensional geometry via a Hewitt\-Manning structural probe\[[19](https://arxiv.org/html/2609.00518#bib.bib21)\]trained to predict squaredCαC\_\{\\alpha\}distances on AB\-Bind and SAbDab structures\[[20](https://arxiv.org/html/2609.00518#bib.bib14)\]\. Finally, developability macro\-Spearmanρ\\rho\(Dev\) captures the recovery of critical biophysical properties for therapeutic applications with a mean\-pooled linear regressor across the five TDC TAP metrics\[[21](https://arxiv.org/html/2609.00518#bib.bib11)\]\. Full details regarding probe architectures, optimization, and data splits are provided in Appendix[C](https://arxiv.org/html/2609.00518#A3)\.

![Refer to caption](https://arxiv.org/html/2609.00518v1/fig2.png)Figure 2:Mask placement distribution influences downstream task performance in antibody MLMs\. a,Performance profiles of eight distinct masking strategies evaluated across predictive tasks spanning five functional axes: paratope prediction \(Para\), contact map prediction \(Cont\), structure probe \(Struc\), developability \(Dev\), and CDR3 infilling \(CDR3\)\. Each strategy excels at a distinct subset of tasks, improving over random masking in at least one area\.b,Empirical probability distribution of mask placements across sequence positions\. Comparing these distributions to performance inademonstrates that strategies excel at tasks functionally aligned with their masking locations\.
## 5Results

Function\-aware masking aligns learned representations with biological priors\.As detailed in Figure[2](https://arxiv.org/html/2609.00518#S4.F2)and Table[1](https://arxiv.org/html/2609.00518#S5.T1), embedding biological priors into the MLM corruption process yields representations that successfully specialize in their targeted functional tasks\. With the exception of structure\-based strategies, every specialist masking strategy achieved the highest score on at least one metric\.

Table 1:Specialist masking strategy results\. Metrics are defined in Section[4](https://arxiv.org/html/2609.00518#S4)and strategies are defined in Section[3\.1](https://arxiv.org/html/2609.00518#S3.SS1)\.untrainedis a randomly initialized control \(no pretraining\)\. Higher is better on all reported metrics\.By directing the masking objective toward residues critical for specific functions, these strategies effectively align the learned representations with their corresponding downstream tasks\. For example,cdrmasking performed best on CDR3 infilling, corroborating previously reported observations\[[5](https://arxiv.org/html/2609.00518#bib.bib1),[6](https://arxiv.org/html/2609.00518#bib.bib16)\]\. Theinterfaceandgermlinemasking strategies performed the best overall, scoring highest across multiple structure\- and paratope\-related tasks\. Interestingly, whileintersectionmasking improved representations for paratope\- and CDR\-related tasks, it largely failed to retain developability and structural gains from theinterfaceandgermlinestrategies it was derived from\. This performance gap likely arises due tointersectionmasking concentrating masks in fewer sequence locations, indicating that maintaining broad competency requires a sufficient diversity of mask locations across the entire sequence\.

structureandstructure\-lrmasking performed worse relative to the other specialists\. While structure\-based masking has demonstrated performance improvements on tasks reliant on structural contacts\[[7](https://arxiv.org/html/2609.00518#bib.bib17)\], antibody frameworks are largely conserved over the pretraining corpus\. The residues most affiliated with modulating structural functions lie in the CDRs, which are targeted more frequently withinterfaceandintersectionmasking; modeling distal contacts in framework regions provides little additional training signal\.

Whilespanmasking does not incorporate a specific biological prior, masking contiguous regions increases the difficulty of the MLM task, a mechanism associated with learning better representations in natural language processing\[[22](https://arxiv.org/html/2609.00518#bib.bib7)\]\. Indeed,spanmasking generalized well across paratope, CDR3, and developability tasks\. However,spanmasking fell short on tasks thatinterfaceandgermlinemasking specialized in, indicating that the placement of masks is more important than the shape for these tasks\.

Overall, by targeting residues most associated with their biological prior, specialist masking strategies effectively align learned representations with their intended downstream tasks\. Furthermore, with the exception ofstructuremasking, every specialist also improved overrandommasking on multiple metrics, includinginterfaceandgermlinemasking, which demonstrated gains across nearly all functional axes\. These results indicate both the upside and necessity of moving away from random masking in antibody MLMs\.

Hybrid masking strategies balance generalization across functional tasks\.While specialist strategies excelled on their targeted objectives, their performance can degrade on unaligned tasks\. To achieve more robust representations, we evaluated seven hybrid masking strategies, each designed to combine multiple biological priors via stochastic mixtures and training curricula\.

Table[2](https://arxiv.org/html/2609.00518#S5.T2)details the performance of each hybrid method\. The best performing hybrid strategies werehybrid\-warmstartandhybrid\-reversemasking\. Interestingly,hybrid\-reversemasking performed best on three of seven tasks despite contradicting the design principle of the other hybrids: masks should be placed uniformly before specializing\. This indicates that important functional axes may be learned early during pretraining, after which specializing provides diminishing returns\.

hybridmasking performed the worst overall, being superseded in performance by both therandomandhybrid\-randommasking controls\. Whilehybrid\-randommasking outperformedhybridmasking by uniformly sampling from specialists, the remaining hybrid strategies surpassed both controls\. This demonstrates that the choice of mixing distribution and learning curriculum is essential for learning representations that generalize across functional axes\. These strategies achieved a higher average task rank \(Figure[3](https://arxiv.org/html/2609.00518#S5.F3)\), particularly on structure\-based tasks\. Alongside their competitive average task rank compared to their specialist counterparts, hybrid strategies also demonstrated a lower overall task variance\. While hybrid strategies seldom performed best on any individual task, their consistent performance across all functional axes indicates their ability to successfully bridge the gap between narrow specialization and broad generalization\.

Table 2:Hybrid strategy results\. Metrics are defined in Section[4](https://arxiv.org/html/2609.00518#S4)and strategies are defined in Section[3\.2](https://arxiv.org/html/2609.00518#S3.SS2)\. Higher is better on all reported metrics\. Each strategy underwent the same training setup with the exception ofhybrid\-warmstart, which was trained for an additional 50,000 steps\.![Refer to caption](https://arxiv.org/html/2609.00518v1/fig3.png)Figure 3:Hybrid masking schemes mitigate task\-specific trade\-offs\.Performance rankings of specialist versus hybrid masking strategies evaluated across multiple tasks\. While specialist strategies can achieve high\-ranking performance on specific tasks, they exhibit significant performance drops on others\. Hybrid strategies mitigate this behavior by combining specialist strategies, compensating for individual task inadequacies to achieve a higher average ranking\.
## 6Conclusion

In this work, we introduced function\-aware masking, a framework that leverages biological priors to impose targeted inductive biases during the pretraining of antibody language models\. While standard random masking treats all residues equally, function\-aware masking aligns the corruption objective directly with the diverse requirements of therapeutic drug discovery\. We demonstrated that specialist masking strategies improve downstream performance on their aligned tasks, achieving up to a14%14\\%improvement on structure\-related tasks and up to a5\.9×5\.9\\timesgain on CDR\-related tasks, all without incurring additional parameter costs at inference\.

As antibody optimization requires balancing multiple complex properties simultaneously, we further developed hybrid masking strategies to mitigate the potential trade\-offs of narrow specialization\. By combining multiple functional priors via stochastic mixtures and training curricula, hybrid models successfully maintain broad generalization and consistently outperform random masking across a diverse array of functional tasks\. Ultimately, our findings reinforce that mask placement acts as an inductive bias over learned representations, offering a flexible, parameter\-free method for training antibody language models\.

Limitations\.Function\-aware masking requires per\-residue labels, which may vary in accuracy or acquisition complexity depending on the strategy\. Evaluations are done on a medium\-sized RoFormer model \(≈\\approx38M params\); further testing would elucidate the impact of masking priors on a larger scale\. Hybrid strategies were manually designed; systematic optimization of mixing distributions may yield additional gains\.

Future work\.Masking strategies represent a promising mechanism for shaping the organization of learned representations during pretraining, yet the extent to which they influence downstream generalization remains poorly understood\. Early evidence indicates that masking can be particularly useful for disentangling higher\-order effects in proteins\[[7](https://arxiv.org/html/2609.00518#bib.bib17)\], a central problem in protein engineering\[[23](https://arxiv.org/html/2609.00518#bib.bib22),[24](https://arxiv.org/html/2609.00518#bib.bib24)\]and variant effect prediction\[[25](https://arxiv.org/html/2609.00518#bib.bib25)\]\. An interesting future direction could incorporate explainability\-driven approaches into the masking strategy\[[26](https://arxiv.org/html/2609.00518#bib.bib23)\], enabling interpretable control over corruption schemes\.

Acknowledgments\.This work was supported in part by the HIVES program at Georgia Tech Research Institute \(GTRI\) and Georgia Institute of Technology start\-up funds\.

## References

- \[1\]J\. Leem, L\. S\. Mitchell, J\. H\. R\. Farmery, J\. Barton, and J\. D\. Galson\(2022\)Deciphering the language of antibodies using self\-supervised learning\.Patterns3\(7\),pp\. 100513\.External Links:ISSN 2666\-3899,[Document](https://dx.doi.org/10.1016/j.patter.2022.100513)Cited by:[§1](https://arxiv.org/html/2609.00518#S1.p1.1)\.
- \[2\]J\. Barton, A\. Gaspariunas, D\. A\. Yadin, J\. Dias, F\. L\. Nice, D\. H\. Minns, O\. Snudden, C\. Povall, S\. Valle Tomas, H\. Dobson, J\. H\. R\. Farmery, J\. Leem, and J\. D\. Galson\(2024\)A generative foundation model for antibody sequence understanding\.bioRxiv\.External Links:[Document](https://dx.doi.org/10.1101/2024.05.22.594943)Cited by:[§1](https://arxiv.org/html/2609.00518#S1.p1.1),[§4](https://arxiv.org/html/2609.00518#S4.p1.1)\.
- \[3\]T\. H\. Olsen, I\. H\. Moal, and C\. M\. Deane\(2022\)AbLang: an antibody language model for completing antibody sequences\.Bioinformatics Advances2\(1\),pp\. vbac046\.External Links:ISSN 2635\-0041,[Document](https://dx.doi.org/10.1093/bioadv/vbac046)Cited by:[§1](https://arxiv.org/html/2609.00518#S1.p1.1)\.
- \[4\]T\. H\. Olsen, I\. H\. Moal, and C\. M\. Deane\(2024\)Addressing the antibody germline bias and its effect on language models for improved antibody design\.Bioinformatics40\(11\),pp\. btae618\.External Links:ISSN 1367\-4811,[Document](https://dx.doi.org/10.1093/bioinformatics/btae618)Cited by:[§1](https://arxiv.org/html/2609.00518#S1.p1.1)\.
- \[5\]K\. Ng and B\. Briney\(2025\)Focused learning by antibody language models using preferential masking of non\-templated regions\.Patterns6\(6\),pp\. 101239\.External Links:ISSN 2666\-3899,[Document](https://dx.doi.org/10.1016/j.patter.2025.101239)Cited by:[§1](https://arxiv.org/html/2609.00518#S1.p2.1),[§2](https://arxiv.org/html/2609.00518#S2.p3.1),[§3\.1](https://arxiv.org/html/2609.00518#S3.SS1.p4.1),[§5](https://arxiv.org/html/2609.00518#S5.p2.1)\.
- \[6\]M\. Talaei, K\. C\. Walker, B\. Hao, E\. Jolley, Y\. Jin, D\. Kozakov, J\. Misasi, S\. Vajda, I\. Ch\. Paschalidis, and D\. Joseph\-McCarthy\(2025\)Preferential CDR masking in paired antibody language models improves binding affinity prediction\.bioRxiv\.External Links:[Document](https://dx.doi.org/10.1101/2025.10.31.685149)Cited by:[§1](https://arxiv.org/html/2609.00518#S1.p2.1),[§2](https://arxiv.org/html/2609.00518#S2.p3.1),[§5](https://arxiv.org/html/2609.00518#S5.p2.1)\.
- \[7\]T\. Walton, A\. Goel, and A\. Aghazadeh\(2026\)Structure\-aware masking for protein representation learning\.arXiv\.External Links:2605\.16581,[Link](https://arxiv.org/abs/2605.16581)Cited by:[§1](https://arxiv.org/html/2609.00518#S1.p2.1),[§2](https://arxiv.org/html/2609.00518#S2.p1.1),[§5](https://arxiv.org/html/2609.00518#S5.p3.1),[§6](https://arxiv.org/html/2609.00518#S6.p4.1)\.
- \[8\]J\. A\. Ruffolo, L\. Chu, S\. P\. Mahajan, and J\. J\. Gray\(2023\)Fast, accurate antibody structure prediction from deep learning on massive set of natural antibodies\.Nature Communications14\(1\),pp\. 2389\.External Links:ISSN 2041\-1723,[Document](https://dx.doi.org/10.1038/s41467-023-38063-x)Cited by:[§B\.1](https://arxiv.org/html/2609.00518#A2.SS1.p1.1),[§1](https://arxiv.org/html/2609.00518#S1.p4.1),[§4](https://arxiv.org/html/2609.00518#S4.p2.1)\.
- \[9\]J\. Devlin, M\. Chang, K\. Lee, and K\. Toutanova\(2019\)BERT: pre\-training of deep bidirectional transformers for language understanding\.InProceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 \(Long and Short Papers\),Minneapolis, Minnesota,pp\. 4171–4186\.External Links:[Document](https://dx.doi.org/10.18653/v1/N19-1423)Cited by:[§2](https://arxiv.org/html/2609.00518#S2.p1.1),[§3](https://arxiv.org/html/2609.00518#S3.p2.1)\.
- \[10\]Z\. Zhang, X\. Han, Z\. Liu, X\. Jiang, M\. Sun, and Q\. Liu\(2019\)ERNIE: enhanced language representation with informative entities\.InProceedings of the 57th Annual Meeting of the Association for Computational Linguistics,Florence, Italy,pp\. 1441–1451\.External Links:[Document](https://dx.doi.org/10.18653/v1/P19-1139)Cited by:[§2](https://arxiv.org/html/2609.00518#S2.p1.1)\.
- \[11\]I\. Yamada, A\. Asai, H\. Shindo, H\. Takeda, and Y\. Matsumoto\(2020\)LUKE: deep contextualized entity representations with entity\-aware self\-attention\.InProceedings of the 2020 Conference on Empirical Methods in Natural Language Processing \(EMNLP\),Online,pp\. 6442–6454\.External Links:[Document](https://dx.doi.org/10.18653/v1/2020.emnlp-main.523)Cited by:[§2](https://arxiv.org/html/2609.00518#S2.p1.1)\.
- \[12\]C\. Lin, T\. Miller, D\. Dligach, S\. Bethard, and G\. Savova\(2021\)EntityBERT: entity\-centric masking strategy for model pretraining for the clinical domain\.InProceedings of the 20th Workshop on Biomedical Language Processing,Online,pp\. 191–201\.External Links:[Document](https://dx.doi.org/10.18653/v1/2021.bionlp-1.21)Cited by:[§2](https://arxiv.org/html/2609.00518#S2.p1.1)\.
- \[13\]D\. Yang, Z\. Zhang, and H\. Zhao\(2023\)Learning better masking for better language model pre\-training\.InProceedings of the 61st Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),Toronto, Canada,pp\. 7255–7267\.External Links:[Document](https://dx.doi.org/10.18653/v1/2023.acl-long.400)Cited by:[§3\.2](https://arxiv.org/html/2609.00518#S3.SS2.p1.1)\.
- \[14\]J\. Su, M\. Ahmed, Y\. Lu, S\. Pan, W\. Bo, and Y\. Liu\(2024\)RoFormer: enhanced transformer with rotary position embedding\.Neurocomputing568,pp\. 127063\.External Links:ISSN 0925\-2312,[Document](https://dx.doi.org/10.1016/j.neucom.2023.127063)Cited by:[§4](https://arxiv.org/html/2609.00518#S4.p1.1)\.
- \[15\]A\. Kovaltsuk, J\. Leem, S\. Kelm, J\. Snowden, C\. M\. Deane, and K\. Krawczyk\(2018\)Observed antibody space: a resource for data mining next\-generation sequencing of antibody repertoires\.The Journal of Immunology201\(8\),pp\. 2502–2509\.External Links:ISSN 1550\-6606,[Document](https://dx.doi.org/10.4049/jimmunol.1800708)Cited by:[§4](https://arxiv.org/html/2609.00518#S4.p2.1)\.
- \[16\]J\. Dunbar and C\. M\. Deane\(2016\)ANARCI: antigen receptor numbering and receptor classification\.Bioinformatics32\(2\),pp\. 298–300\.External Links:ISSN 1367\-4811,[Document](https://dx.doi.org/10.1093/bioinformatics/btv552)Cited by:[§4](https://arxiv.org/html/2609.00518#S4.p2.1)\.
- \[17\]E\. Liberis, P\. Veličković, P\. Sormanni, M\. Vendruscolo, and P\. Liò\(2018\)Parapred: antibody paratope prediction using convolutional and recurrent neural networks\.Bioinformatics34\(17\),pp\. 2944–2950\.External Links:ISSN 1367\-4811,[Document](https://dx.doi.org/10.1093/bioinformatics/bty305)Cited by:[§4](https://arxiv.org/html/2609.00518#S4.p2.1)\.
- \[18\]J\. Dunbar, K\. Krawczyk, J\. Leem, T\. Baker, A\. Fuchs, G\. Georges, J\. Shi, and C\. M\. Deane\(2014\)SAbDab: the structural antibody database\.Nucleic Acids Research42\(D1\),pp\. D1140–D1146\.External Links:ISSN 1362\-4962,[Document](https://dx.doi.org/10.1093/nar/gkt1043)Cited by:[§4](https://arxiv.org/html/2609.00518#S4.p2.1)\.
- \[19\]J\. Hewitt and C\. D\. Manning\(2019\)A structural probe for finding syntax in word representations\.InProceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 \(Long and Short Papers\),Minneapolis, Minnesota,pp\. 4129–4138\.External Links:[Document](https://dx.doi.org/10.18653/v1/N19-1419)Cited by:[§C\.1](https://arxiv.org/html/2609.00518#A3.SS1.p4.1),[§4](https://arxiv.org/html/2609.00518#S4.p4.1)\.
- \[20\]S\. Sirin, J\. R\. Apgar, E\. M\. Bennett, and A\. E\. Keating\(2016\)AB\-Bind: antibody binding mutational database for computational affinity predictions\.Protein Science25\(2\),pp\. 393–409\.External Links:ISSN 1469\-896X,[Document](https://dx.doi.org/10.1002/pro.2829)Cited by:[§4](https://arxiv.org/html/2609.00518#S4.p4.1)\.
- \[21\]M\. I\. J\. Raybould, C\. Marks, K\. Krawczyk, B\. Taddese, J\. Nowak, A\. P\. Lewis, A\. Bujotzek, J\. Shi, and C\. M\. Deane\(2019\)Five computational developability guidelines for therapeutic antibody profiling\.Proceedings of the National Academy of Sciences116\(10\),pp\. 4025–4030\.External Links:ISSN 1091\-6490,[Document](https://dx.doi.org/10.1073/pnas.1810576116)Cited by:[§4](https://arxiv.org/html/2609.00518#S4.p4.1)\.
- \[22\]M\. Joshi, D\. Chen, Y\. Liu, D\. S\. Weld, L\. Zettlemoyer, and O\. Levy\(2020\)SpanBERT: improving pre\-training by representing and predicting spans\.Transactions of the Association for Computational Linguistics8,pp\. 64–77\.External Links:ISSN 2307\-387X,[Document](https://dx.doi.org/10.1162/tacl%5Fa%5F00300)Cited by:[§5](https://arxiv.org/html/2609.00518#S5.p4.1)\.
- \[23\]D\. Tsui and A\. Aghazadeh\(2024\)On recovering higher\-order interactions from protein language models\.External Links:2405\.06645,[Link](https://arxiv.org/abs/2405.06645)Cited by:[§6](https://arxiv.org/html/2609.00518#S6.p4.1)\.
- \[24\]D\. Tsui, K\. Talreja, and A\. Aghazadeh\(2025\)Sparse autoencoders for low\-NNprotein function prediction and design\.External Links:2508\.18567,[Link](https://arxiv.org/abs/2508.18567)Cited by:[§6](https://arxiv.org/html/2609.00518#S6.p4.1)\.
- \[25\]T\. Walton, D\. Tsui, L\. Fogel, R\. Chagas, D\. Huard, R\. Lieberman, and A\. Aghazadeh\(2025\)GOLF: a generative ai framework for pathogenicity prediction of myocilin olf variants\.InProceedings of the 20th Machine Learning in Computational Biology meeting,Proceedings of Machine Learning Research, Vol\.311,pp\. 148–161\.External Links:[Link](https://proceedings.mlr.press/v311/walton25a.html)Cited by:[§6](https://arxiv.org/html/2609.00518#S6.p4.1)\.
- \[26\]D\. Tsui, A\. Musharaf, Y\. E\. Erginbas, J\. Kang, and A\. Aghazadeh\(2025\)SHAP zero explains biological sequence models with near\-zero marginal cost for future queries\.InAdvances in Neural Information Processing Systems,Vol\.38,pp\. 83840–83891\.Cited by:[§6](https://arxiv.org/html/2609.00518#S6.p4.1)\.

## Appendix

## Appendix AAdditional Results

We present the full results table in this section, combining experiments from Table[1](https://arxiv.org/html/2609.00518#S5.T1)and Table[2](https://arxiv.org/html/2609.00518#S5.T2)\.

Table 3:Full results table across all strategies presented in the main text\.
## Appendix BImplementation Details

This section outlines implementation details for specialist and hybrid masking schemes presented in the main text\. Furthermore, we visualize empirical probability distributions in Figure[4](https://arxiv.org/html/2609.00518#A2.F4)for thirteen methods tested in this work \(excludinghybridandhybrid\-randommasking\) with CDRs highlighted for reference\.

![Refer to caption](https://arxiv.org/html/2609.00518v1/fig_app_strategies.png)Figure 4:Empirical per\-position mask probabilities for all thirteen strategies evaluated in this work\.### B\.1Specialists

Structure and structure\-LR\.Structures for each antibody are predicted using IgFold\[[8](https://arxiv.org/html/2609.00518#bib.bib12)\], an antibody\-specific structure predictor\. We build a per\-residuekk\-nearest neighbors graph by computing EuclideanCα−CαC\_\{\\alpha\}\-C\_\{\\alpha\}distance\.structureutilizesk=5k=5nearest neighbors, masking residues which form structural neighborhoods\. From these neighborhoods, only a few masks are sampled\. This is done to provide enough structural context to decode the masked residue\. Structure\-LR follows the same logic, but restricts residue neighborhoods to only include corresponding sequence locations given that they are more than four indices apart\.

### B\.2Hybrid Masking Strategies

Hybrid masking strategies sample masks from a mixing distribution specified by a timesteptt, combining specialist strategies as a stochastic mixture\. We describe these learning curricula in this section, indicating the mixing distribution for each training step interval\.

Hybrid \(default\)\.The baseline curriculum for the remaining hybrid strategies\. Starts with higher emphasis onrandomandspanmasking, gradually transitioning to more weight oncdrandinterface\. After step 40,000, the mixing strategy covers each strategy more generally\.

Table 4:Mixing schedule for hybrid \(default\)\.Hybrid\-stretched\.Begins with more probability mass assigned torandomandspan, and gradually stretches out the same strategy ashybridover more steps\.

Table 5:Mixing schedule for hybrid\-stretched\.Hybrid\-reverse\.Same strategy ashybrid\-stretchedbut in reverse, stretched over more steps\.

Table 6:Mixing schedule for hybrid\-reverse\.Hybrid\-weighted\.Focuses more heavily on the best performing specialist,interface\.

Table 7:Mixing schedule for hybrid\-weighted\.Hybrid\-perbatch\.Same probability distribution ashybrid, but instead of sampling the masking strategy for each sequence, samples a fixed strategy per batch\. Less computational overhead than the other hybrid strategies\.

Hybrid\-warmstart\.Continues training of aninterfacemasking model for an additional 50,000 steps\. Samples from the following mixing distribution at each step:\[0\.10,0\.20,0\.10,0\.15,0\.20,0\.25\]\[0\.10,0\.20,0\.10,0\.15,0\.20,0\.25\]\.

## Appendix CExperimental Setup

Each from\-scratch run takes≈\\approx5 hours on one RTX A6000 \(48 GiB\) at 7\.2 it/s under fp16;hybrid\-warmstartadds≈\\approx2 hours; one\-time IgFold prediction cost≈\\approx240 GPU\-hours on 8 GPUs in parallel\.

### C\.1Probe Details

All evaluation metrics in Section[4](https://arxiv.org/html/2609.00518#S4), with the exception of zero\-shot CDR3 infilling, are obtained by linear probing on frozen representations across four probing tasks: paratope, contact map, structure, and developability\. For each task, the pretrained encoder is fixed and its sequence representations are extracted from the last layer\. A task\-specific head is then trained on top of this representation; the encoder never receives any gradient updates\. Heads are optimized with AdamW \(weight decay0\.010\.01\) at a learning rate of10−310^\{\-3\}under a linear warmup \(0\.10\.1\) with a cosine schedule, gradient norm clipping set at1\.01\.0, at full precision\. For each task we trained multiple independent probes across at least three seeds\. At every epoch, the head is scored on the held\-out validation split using that task’s early stopping metric \(Table[8](https://arxiv.org/html/2609.00518#A3.T8)\)\. The best validation checkpoint is retained and used for the test set evaluation\. We additionally fit a decision threshold on the binary paratope task validation by maximizing Youden’s J for paratope MCC; AUPRC and AUROC are threshold\-free\. The seeds randomize head initialization, dropout, and minibatch order only: the data split and embeddings are identical across seeds\.

Each probing task has its own fixed split\. The paratope and developability probes utilize the default splits provided by TDC\. TDC SAbDab\_Liberis is partitioned at random into716/102/205716/102/205antibody sequences \(70/10/20\), and TAP is partitioned per property and merged by antibody identifier, yielding 241 sequences \(70/10/20\)\. The two structural probes are split by PDB entry at 60/20/20 such that heavy and light chains belonging to the same complex are never separated across splits\. The contact map probe utilizes 487 SAbDab crystal structures, assigned292/97/98292/97/98by entry, yielding367/121/130367/121/130chains\. The structure probe merges AB\-Bind and SAbDab into a single pool, deduplicated by entry identifier \(487 from SAbDab, 31 from AB\-Bind\)\. The final structure split is309/103/104309/103/104by entry and407/134/139407/134/139by chain\. All splits are deterministic and identical across every masking strategy and pretraining seed\. CDR3 infilling is performed over1,0001,000sequences from the pretraining hold\-out set\.

The paratope teacher is a per\-token classifier fine\-tuned on top of the pretrained AntiBERTa2 encoder using the same TDC SAbDab\_Liberis partition as the paratope probe \(716/102/205716/102/205\)\. It is fit on the training split only, early\-stopped on validation AUPRC, and reported on the test split\. The encoder and classification head are optimized jointly with AdamW learning rates of2×10−52\\times 10^\{\-5\}and5×10−45\\times 10^\{\-4\}respectively, batch size 16, cosine schedule with10%10\\%warmup, under class\-weighted binary cross\-entropy\. As the teacher and probe call the same split function, no test antibody was labeled by the teacher during teacher training; teacher predictions are computed only over the pretraining set, never over evaluation sequences\. No structure teacher is trained in this work\. Instead, structures are predicted with IgFold, a pretrained antibody structure predictor\. To validate the accuracy of the predictions, we compared predicted neighborhoods against real SAbDab crystal structures\. Across 316 heavy chains, a median of93%93\\%of each residue’s five predicted nearestCαC\_\{\\alpha\}neighbors also appear among its five nearest neighbors in the crystal structure \(4%4\\%for randomly chosen neighbors\), with a median FvCαC\_\{\\alpha\}RMSD of0\.920\.92Å\. IgFold is applied only to pretraining sequences to determine mask placement; any overlap between its training data and our structural benchmarks cannot transmit evaluation labels into the encoder\.

Probe heads\.Paratope \(Para\): a per\-token linear classifier trained with class\-weighted binary cross\-entropy, scored by AUPRC and MCC\. Contact map \(Cont\): a bilinear form over residue pair embeddings \(upper triangle, amino acid positions only\) trained with binary CE, scored by AUROC and long\-range precision\-at\-LL\. Structure probe \(Struct\): a Hewitt\-Manning linear structural probe\[[19](https://arxiv.org/html/2609.00518#bib.bib21)\]that predicts squaredCαC\_\{\\alpha\}distances, trained with masked MSE and scored by the Spearmanρ\\rhobetween predicted and true distances\. Developability \(Dev\): mean pooling over residue tokens followed by a linear regressor trained with MSE onzz\-scored targets, scored by the macro\-averaged Spearmanρ\\rhoover five TAP properties \(encoder hidden size is512512dimensions\)\.

Table 8:Linear\-probe configurations\. Patience00disables early stopping \(developability trains its full budget\)\.