通过功能感知掩码学习任务特定的抗体表示

# 基于函数感知掩码学习任务特异性抗体表征
来源：https://arxiv.org/html/2609.00518
Thomas A. Walton  
Amirali Aghazadeh  
*共同第一作者  
单位：佐治亚理工学院电气与计算机工程学院  
通讯作者：[email protected]  
单位：佐治亚理工学院计算机科学学院  

###### 摘要

通过掩码语言建模（MLM）预训练的抗体特异性语言模型，学习了对下游序列设计和属性预测任务至关重要的表征。然而，掩码过程本身很少在预训练中被用作归纳偏置的来源。虽然优先掩码互补决定区（CDR）可改善结合相关预测，但抗体在多种功能上具有多样的生物学先验。本文引入函数感知掩码，这是一族预训练算法，将掩码放置与特定功能先验（例如来自IMGT注释或结构预测）对齐，以塑造学习到的表征空间。我们表明，这些专用掩码策略显著提升了其相应目标的表现，在结构相关任务上最高提升14%，在CDR相关任务上最高提升5.9倍。为进一步提升多个功能轴上的性能，我们开发了混合掩码策略，整合多种先验，平衡结合、结构和生物物理目标上的重建。我们的结果表明，基于信息的掩码放置为在抗体语言模型训练中施加功能性归纳偏置提供了一种无需额外参数的机制。

## 1 引言

治疗性抗体的开发处于药物发现的前沿，为复杂的肿瘤学、自身免疫和感染性疾病提供靶向治疗。学习有表现力的序列表征对于优化这些生物制剂以实现临床可行性至关重要。在掩码语言建模（MLM）目标下训练的抗体特异性语言模型（abLMs），如AntiBERTa[1 (https://arxiv.org/html/2609.00518#bib.bib2), 2 (https://arxiv.org/html/2609.00518#bib.bib3)] 和 AbLang[3 (https://arxiv.org/html/2609.00518#bib.bib4), 4 (https://arxiv.org/html/2609.00518#bib.bib5)]，其编码的表征在多样化的序列设计和优化任务中展现出强大的泛化能力。

这些模型的改进主要源于预训练数据的扩展和架构的优化。然而，近期工作表明，掩码过程本身作为一种归纳偏置，塑造了所学表征在任务间的泛化方式[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)]。具体而言，优先掩码抗体互补决定区（CDR）可提升CDR填充、结合亲和力和结合特异性任务的迁移性能。

然而，CDR掩码仅针对抗体功能的一个方面。治疗性抗体还必须稳定折叠、抵抗聚集，并表现出突变稳健性以确保临床可开发性。控制这些特性的残基通常位于CDR之外；它们分布在提供结构支架的框架区。CDR掩码及其变体的成功指向了一个更广泛的原则，并提出了一个更普遍的问题：是否可以使用任何任意的功能先验来设计抗体掩码策略，以使表征与相应的下游任务对齐？

参考图注 图1：函数感知掩码将表征学习集中在关键的抗体区域。a，抗体-抗原复合物（PDB ID: 3HFM）的刚性结构，说明功能变异主要集中在一个稀疏的序列位置子集。b，通过腐败（掩码）和重建目标，在数十万抗体上训练的掩码语言模型（MLM）隐式地学习了序列变异如何影响功能。c，随机掩码将所有序列位置视为同等重要；然而，对下游适应度最具决定性的序列位置集中在少数区域。函数感知掩码利用功能注释指导掩码放置，将训练信号聚焦于这些相关区域。
为解决此问题，我们引入了函数感知掩码（图1 (https://arxiv.org/html/2609.00518#S1.F1)），这是一族掩码策略，将掩码放置偏向与特定功能任务最相关的残基。每种策略使用从序列标签或模型预测（例如IMGT标签或IgFold预测[8 (https://arxiv.org/html/2609.00518#bib.bib12)]）推导出的功能信号对残基进行评分，并在训练过程中优先掩码这些位置（图1 (https://arxiv.org/html/2609.00518#S1.F1)c）。重要的是，函数感知掩码仅塑造训练目标：所得模型在推理时仍仅依赖序列，且不产生额外的参数成本。这些掩码策略（下文称为专用策略）在与其先验对齐的任务上优于随机掩码。为在多个功能轴上同时持续提升性能，我们进一步引入了混合掩码策略，通过随机混合和学习课程整合功能先验。我们在预训练设置中评估了函数感知掩码，并在涵盖结合、结构、可开发性和生成目标的五个下游任务上评估了性能。我们的贡献如下：

1. 我们引入了函数感知抗体掩码，一类根据下游任务施加功能性归纳偏置的掩码算法。我们的实现在此开源：https://github.com/amirgroup-codes/function-aware-masking.git
2. 我们证明，专注于单一生物学先验的专用掩码策略，在与其先验对齐的任务上优于随机掩码。专用策略在结构相关任务上性能提升高达14%，在CDR相关任务上提升高达5.9倍。
3. 为进一步提升泛化能力，我们证明混合掩码策略提高了多个功能轴上的表征质量，缓解了专用策略在未对齐任务上的性能下降。

## 2 背景

腐败是MLM训练的一个关键组成部分。最常见的腐败方案是用掩码标记 `[MASK]` 替换标记，历史上通常以固定的概率（通常为15%）在序列位置上均匀随机应用[9 (https://arxiv.org/html/2609.00518#bib.bib6)]。自然语言处理领域的基础性工作表明，掩码放置可以影响学习到的表征，使其为推断语义结构相关的任务做好准备[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)]。我们最近将这一思想扩展到蛋白质序列，通过基于结构接触指导掩码放置，发现所得表征在一系列外推任务上表现出优越的泛化能力[7 (https://arxiv.org/html/2609.00518#bib.bib17)]。

抗体中的生物学先验。虽然一般蛋白质受益于结构先验，但抗体具有高度特化的结构，由多个不同的功能驱动因子控制。抗原识别主要由六个CDR介导，形成连续的、高度可变的区域。在这些环内，一个特定的残基子集（互补位）与目标形成直接界面。在亲和力成熟过程中，这些结合界面及其相应的框架区经历体细胞超突变（SHM），相对于遗传的生殖系序列引入变异。除了结合之外，抗体的整体稳定折叠由结构接触网络维持，包括形成长程相互作用、构成表观遗传中心并决定构象的远程相互作用。值得注意的是，上述功能驱动因子仅占抗体序列的一小部分。由此可见，在MLM训练过程中，应用于这些功能关键区域之外的大部分腐败，在给定周围上下文的情况下很容易解码。

抗体掩码。先前的工作已经探索了抗体的掩码策略，主要集中在CDR上。Ng & Briney[5 (https://arxiv.org/html/2609.00518#bib.bib1)] 引入了对非模板化CDR3的优先掩码，发现所学习到的表征更好地泛化到CDR3残基恢复，并改善了结合特异性预测。Talaei等人[6 (https://arxiv.org/html/2609.00518#bib.bib16)] 通过引入混合的CDR和框架掩码策略扩展了这项工作，进一步证明了对结合亲和力预测的改进。以前的工作主要集中在CDR上，而这项工作探索了与前面讨论的生物学先验相一致的、针对更大范围功能相关任务的掩码策略。

## 3 方法

我们引入函数感知掩码，一套利用抗体生物学先验以在多种任务上实现更优表征的掩码算法。函数感知掩码分为两类：专注于单一任务的专用策略，以及融合多种任务先验的混合策略。

掩码公式化与预计算先验。为了实现有针对性的掩码采样同时限制计算开销，我们预先计算了每个残基的生物学标签，并将其与训练数据一起缓存。标签包括CDR位置、互补位概率、生殖系突变状态和结构拓扑。在数据整理过程中，为序列S（不包括 `[CLS]/[SEP]` 标记）采样一个掩码集M。为了隔离任何给定策略相对于其他策略的效果，我们将掩码率固定为0.15。然后，选定的位置经历BERT引入的标准MLM腐败过程[9 (https://arxiv.org/html/2609.00518#bib.bib6)]：80%的标记被替换为 `[MASK]` 标记，10%被替换为随机氨基酸，剩余的10%保持不变。

### 3.1 专用策略

专用掩码策略的目标是将抗体MLM学习到的表征与特定的下游任务或功能属性对齐。对于基于权重的策略，掩码采样器为每个位置分配一个权重w_i，编码其与目标属性的相关性。权重然后被归一化为针对预期预算15%校准的每个位置的掩码概率（公式1 (https://arxiv.org/html/2609.00518#S3.E1)）。掩码位置由独立的伯努利抽样确定（公式2 (https://arxiv.org/html/2609.00518#S3.E2)）。对于依赖先前放置掩码的非基于权重的策略（span, structure, structure-lr），掩码被迭代放置直到达到15%的固定预算。

公式1:  
p_i = min(0.15 * w_i / \bar{w}, 1) * 1[i ∈ S], \quad \bar{w} = (1/|S|) ∑_{j∈S} w_j, \quad E[|M|] = 0.15 * |S|.

公式2:  
m_i ~ Bernoulli(p_i) 独立, \quad M = { i: m_i = 1 }.

在此公式化下，我们评估了七种旨在针对不同功能、进化和结构先验的专用策略。具体实现的细节可在第4节 (https://arxiv.org/html/2609.00518#S4) 找到。

cdr：通过过采样互补决定环来针对高变区（w_FW:w_CDR1:w_CDR2:w_CDR3 = 1:3:3:6），将CDR3掩码推广到其余CDR[5 (https://arxiv.org/html/2609.00518#bib.bib1)]。

span：与拓扑无关，使用几何分布（p=0.2, l_max=10）掩码连续的序列位置片段。

interface：通过将权重比6:1放置在互补位接触残基上，针对物理结合界面。

germline：通过将权重比6:1放置在生殖系突变残基上，针对抗体的进化轨迹。

intersection：一种联合专用策略，将互补位权重和生殖系权重相乘：w_i = w_i^{para} * w_i^{germ}。该策略针对负责亲和力成熟的残基，将掩码集中在既是抗原接触位点又是来自生殖系体细胞突变的位置。

structure：一种3D感知方案，基于预测的空间邻近性对残基进行加权。利用结构的五个最近邻来掩码结构邻域。掩码在邻域之间采样，留下足够的空间上下文来解码被掩码的残基。

structure-lr：structure的扩展，应用更严格的分区规则：仅考虑序列索引间隔大于4的结构耦合。优先在重建过程中建模长程接触。

### 3.2 混合策略

专用策略针对一种期望的抗体属性；然而，可能有许多这样的属性需要优化。为解决此问题，我们开发了混合策略，在训练过程中随机混合功能先验。给定一组K个掩码策略，我们定义混合分布 π(t) = (p_1(t), ..., p_K(t))，其中p_i(t)表示在训练步t抽样专用策略i的概率（∑_s p_s(t) = 1）。混合分布遵循两个原则：掩码应在专业化前均匀采样[13 (https://arxiv.org/html/2609.00518#bib.bib26)]，并且表现最佳的专用策略应有最高的被采样概率。采样通过抽取策略i ~ π(t)并按序列应用掩码进行。我们在训练过程中动态调整π(t)，以偏向于抽取哪些功能先验。混合分布构建的细节可在附录B.2 (https://arxiv.org/html/2609.00518#A2.SS2) 找到。

除非另有说明，混合策略在以下专用策略之间插值，引用顺序为：[random, cdr, span, structure-lr, interface, germline]。我们评估以下混合策略：

hybrid：以下混合策略的基础。强调random和span掩码，然后逐渐转向专用策略。

hybrid-random：一种随机化的混合控制策略，从π(t)中均匀采样。

hybrid-stretched：开始时将混合分布严重偏向random和span。以与hybrid相同的分布结束。将混合中的转换延迟到更多时间步骤。

hybrid-reverse：反转hybrid-stretched的应用方向。

hybrid-weighted：以类似于hybrid-stretched的通用策略开始，将重点转向interface掩码，然后是germline和interface掩码的混合，最后以更平衡的分布结束。

hybrid-perbatch：与hybrid相同的计划，但按序列批次抽取单个专用策略，而不是按序列抽取。

hybrid-warmstart：持续预训练而非课程学习。采用一个完全训练好的interface掩码模型并