通过分类器引导的胚系吸收离散扩散实现抗体序列的条件生成

# 利用分类器引导的祖本吸收离散扩散实现抗体序列的条件生成

来源：https://arxiv.org/html/2605.06720

Justin Sanders¹ jsander1@cs\.washington\.edu
Luca Giancardo² lugianca@amazon\.com
Lan Guo² languo@amazon\.com
Yue Zhao² yuezhaom@amazon\.com
Kemal Sonmez² ksonmez@amazon\.com
Nina Cheng² cxinyun@amazon\.com
Melih Yilmaz² melihyz@amazon\.com

¹ 华盛顿大学 Paul G. Allen 计算机科学与工程学院
² 亚马逊网络服务生命科学部门

###### 摘要

抗体疗法是现代医学中最为成功的药物之一，然而，计算设计具有理想结合性和可开发性特性的抗体仍然面临挑战。虽然蛋白质语言模型（pLMs）已成为抗体序列设计的强大工具，但现有方法主要存在两个关键局限性：它们主要记忆祖本（germline）序列，而非对具有生物学意义的体细胞变异进行建模；此外，它们对灵活分类器引导的条件生成支持有限。我们通过两项主要贡献解决了这些挑战。首先，我们证明离散扩散微调在抗体序列上实现了强大的语言建模性能，同时允许基于任何现成分类器进行条件生成。其次，我们引入了祖本吸收扩散（germline absorbing diffusion），这是一种对离散扩散噪声过程的新颖修改，其中祖本序列——而非掩码序列——作为吸收态。这种基于生物学的归纳偏置限制了模型学习从祖本到观测序列的轨迹，有效地将遗传变异和 V(D)J 重排统计从学习到的分布中排除，从而大幅减轻了祖本偏差。我们表明，祖本扩散将非祖本残基预测准确率从 26% 提高到了 46%，接近由真实生物学变异性设定的理论上限。随后，我们在改善疏水性和预测结合亲和力的抗体采样任务中展示了祖本扩散模型的效用。在这两项任务中，我们的模型在类别遵循度和样本质量之间显示出更好的权衡，显著优于 EvoProtGrad（一种流行的基于梯度离散马尔可夫链蒙特卡洛从 pLMs 中采样的策略）。值得注意的是，在一项优化临床抗体 Emibetuzumab 对其抗原 HGFR 结合亲和力的实验中，我们的模型生成的设计平均具有高出 48% 的 Boltz-2 预测抗体-抗原复合物得分，并产生了所有设计中预测结合能力最强的分子。

## 1 引言

抗体是现代医学的基石，代表最具成效的治疗分子类别之一 Mullard (2021); Liu et al. (2020)。它们的高结合特异性以及与人类免疫系统的兼容性使其成为治疗癌症、自身免疫疾病和感染性病原体的理想候选者。然而，发现和优化具有合适结合亲和力、特异性和可开发性特性组合的抗体仍然是一个昂贵且耗时的过程，这促使开发计算方法以加速设计周期。

计算抗体设计历史上依赖于基于物理的能量函数和基于结构的方法 Adolf-Bryfogle et al. (2018); Sivasubramanian et al. (2009)，尽管这些方法强大，但受限于对高质量抗原结构的需求以及穷举构象搜索 prohibitive 的计算成本。蛋白质语言模型（pLMs）的兴起开启了一种互补的以序列为中心的范式，通过在数亿蛋白质序列上进行大规模预训练，获得隐含编码丰富结构和进化信息的表示 Rives et al. (2021); Lin et al. (2023)。抗体特异性 pLMs，如 AntiBERTy Ruffolo et al. (2021)、AbLang Olsen et al. (2022a) 和 IgLM Shuai et al. (2023)，在包括责任预测（liability prediction）、人源化评分和互补决定区（CDR）填充在内的任务中展示了强劲性能。然而，通过在大规模 B 细胞库上使用标准的无监督语言建模设置进行训练，这些模型往往捕获了限制其有效性的不良偏差。

天然 B 细胞序列受到许多卷积生物过程的影响：V、D 和 J 基因片段的遗传变异、发育过程中这些片段的随机重排（V(D)J 重排）、正向选择和亲和力成熟以改善结合，以及中枢/外周耐受以避免自身反应性。虽然其中一些过程对于模型捕获至关重要，但其他过程对于感兴趣的任务则是多余的。许多抗体 pLMs 已被证明主要学习了常见的遗传变异和 V(D)J 重排模式 Olsen et al. (2024); Matsen IV et al. (2025)。因此，它们相对于通用 pLMs 的明显性能提升很大程度上源于记忆祖本，而非准确区分功能性抗体的生物学有意义变异——这对*从头*设计和定向进化运动是一个关键限制。

除了祖本偏差问题外，现有的抗体生成模型，特别是那些允许由另一个模型（例如，具有可开发性或结合亲和力预测器的分类器引导生成）引导条件生成的模型，主要在自回归、单步掩码预测或梯度引导的马尔可夫链蒙特卡洛（MCMC）框架下运行 Madani et al. (2023); Shuai et al. (2023); Emami et al. (2023)。虽然自回归模型携带适用于自然语言的归纳偏置，但由于生物序列不一定具有相同的上下文顺序依赖性，它们可能不太适合抗体。此外，更重要的是，这些方法只能由相对于输入序列完全可微的分类器进行引导，严重限制了可用于引导的模型类型。此外，即使在可计算的情况下，在离散域中工作时，这些梯度可能也不具有语义意义。

> **图 1：** 具有祖本吸收态的离散扩散抗体蛋白质语言模型
> (A) 我们的模型使用分数熵离散扩散（SEDD）框架训练，旨在将祖本序列去噪为观测到的抗体序列。
> (B) 训练完成后，我们的模型允许基于任意分类器进行抗体序列的*从头*生成和定向进化。

离散扩散模型提供了一个引人注目的替代方案 Austin et al. (2021); Lou et al. (2024)。通过将生成框定为离散序列标记上噪声过程的逐渐逆转，它们自然地支持灵活的条件化、高效的填充以及直截了当的分类器引导生成。前向扩散的归纳偏置非常适合受自然选择影响的生物序列。离散扩散模型将序列视为生成过程的结果，其中序列经历随机替换，偏向于高似然区域。这类似于个体突变通过自然选择逐渐增加适应度。

在这项工作中，我们提出了一种新颖的离散扩散抗体 pLM 来解决这些机遇。我们的模型建立在分数熵离散扩散（SEDD）框架 Lou et al. (2024) 之上，并利用 ESM-2 变压器架构 Lin et al. (2023)，从一般蛋白质序列的预训练权重初始化，并在来自观测抗体空间（OAS）数据库 Kovaltsuk et al. (2018); Olsen et al. (2022b) 的超过 2500 万个多样 B 细胞受体序列上使用扩散目标进行微调。为了减轻祖本偏差问题，我们采用了一种新颖的离散扩散框架，其中祖本序列用作噪声过程的吸收态。通过以这种方式显式地以祖本为条件，我们假设模型将更多地致力于学习 B 细胞受体序列分布中与治疗设计最相关的方面。此外，通过更忠实地表示数据背后的真实生成过程，我们为模型引入了有利的归纳偏置。在语言建模指标上评估我们的模型时，我们发现扩散模型非常适合用作抗体 pLMs，且我们的祖本吸收扩散设置显著减少了祖本偏差。在下游*计算*条件生成任务中，我们称之为 MochiDiff 的新模型成功设计了具有理想可开发性和结合特性的序列。

## 2 背景和相关工作

##### 抗体设计

抗体是蛋白质，其靶标特异性主要由六个 CDRs 决定，其中 CDR-H3 是最多样且功能关键的 Chothia et al. (1989); North et al. (2011)。抗体可以用描述氨基酸的字母*序列*表示，和/或用描述其物理*结构*的 3D 原子位置列表表示。传统的计算设计依赖于结构空间中的基于物理的能量函数 Adolf-Bryfogle et al. (2018); Sivasubramanian et al. (2009)，这些方法计算昂贵且需要高质量的抗原结构。深度学习已通过最近的基于扩散的方法（如 DiffAb Luo et al. (2022)、AbDiffuser Martinkus et al. (2023)、RFdiffusion Watson et al. (2023) 和 Chroma Ingraham et al. (2023)）大幅放松了这些限制，实现了无条件及条件引导的主链生成。尽管取得了这些进展，解析抗原结构的可用性有限，以及原子位置高维空间固有的计算成本仍然限制基于结构的设计，从而促使仅能从序列数据操作的 pLM 方法。

##### 蛋白质语言模型与条件生成

在蛋白质序列数据库上的大规模预训练产生了 pLMs，它们在学习的序列表示中编码了丰富的进化、结构和功能信息 Rives et al. (2021); Elnaggar et al. (2021)。ESM-2 Lin et al. (2023) 在数亿序列上使用掩码语言建模目标进行训练，产生残基级嵌入，从中可以通过线性探测恢复二级结构、溶剂可及性和接触图。这使得 pLMs 不仅作为编码器，而且作为条件序列设计的生成先验具有吸引力，出现了两种主要范式。

第一种范式在特定领域的抗体数据上微调 pLMs，以捕获区分抗体与通用蛋白质的免疫球蛋白特异性统计信息。AntiBERTy Ruffolo et al. (2021) 和 AbLang Olsen et al. (2022a) 是 BERT 风格 Devlin et al. (2018) 的掩码语言模型，仅在抗体库序列上预训练，与通用 pLMs 相比，在责任预测、人源化评分和 CDR 填充方面表现出优越性能。然而，AbLang-2 Olsen et al. (2024) 随后表明，这种提升很大程度上源于记忆常见的祖本序列和 V(D)J 重排模式，而非准确区分功能性抗体的生物学有意义的体细胞变异，这对于新型治疗序列的条件生成是一个关键限制。

一些抗体 pLMs 通过在训练期间前缀离散属性或注释令牌来支持条件生成。ProGen Madani et al. (2023) 和 ProGen2 Nijkamp et al. (2023) 在功能家族注释上条件化自回归模型，以实现对广泛蛋白质家族的可控生成。在抗体领域，IgLM Shuai et al. (2023) 采用相同的因果语言建模框架，条件化于链类型和物种令牌，实现靶向抗体序列生成。虽然有效，但架构条件化要求所需属性在训练时预先指定，并伴随足够的标记训练数据，这些条件对于治疗设计核心的许多生物物理特性很少成立。关键在于，由于条件化在训练时被锁定，这些模型中没有一个原生支持灵活的分类器引导生成。

第二种范式利用 pLM 的学习似然景观来事后引导生成，而不修改模型本身。EvoProtGrad Emami et al. (2023) 提供了一个可组合的基于梯度的 MCMC 框架，将一个或多个适应度预测器与 pLM 伪对数似然分数相结合，以引导离散序列优化。这允许任何可微分类器在不重新训练 pLM 的情况下插入。然而，离散序列上的基于梯度的 MCMC 需要引入误差的连续松弛，且对可微引导分类器的要求严重限制了开箱即用可使用的模型。

##### 离散扩散

扩散概率模型 Sohl-Dickstein et al. (2015); Ho et al. (2020); Song et al. (2021) 在连续域（如图像 Rombach et al. (2022)、音频 Kong et al. (2021) 和蛋白质结构生成 Jumper et al. (2021); Yi et al. (2023)）中取得了显著成功；将其扩展到离散序列需要根本不同的前向过程公式。Austin 等人介绍了 D3PM Austin et al. (2021)，这是一个离散扩散的统一框架，概括了多项式 c