INTORDUCTION

背景与动机

蛋白质是由氨基酸线性排列并折叠成三维结构的生物大分子，在生物学中起着关键作用，如调控转录、翻译、信号传导和细胞周期控制等。近年来，基于数据驱动的深度学习方法在蛋白质理解与设计中展现出巨大的潜力，逐渐推动了蛋白质研究从传统物理方法向数据驱动的方向转变。

蛋白质序列与自然语言之间的类比已经被广泛接受，并受到自然语言处理（NLP）领域中语言模型的成功启发。借助Transformer的可扩展性和大规模数据，研究者们开始使用语言模型（如BERT、GPT等）来学习蛋白质序列，取得了显著的成果。这些模型不仅能完成预测任务（如功能属性探测、从单序列预测蛋白质结构等），还能生成新的蛋白质序列。然而，现有的蛋白质语言模型在生成和理解方面还未达到其最大潜力。

现有问题

  1. Masked 语言模型的局限：
    • 类似BERT的掩码语言模型（Masked-LMs），如ESM系列，在蛋白质预测任务中表现出色，但由于双向感知的特性，使得其难以进行序列生成。此外，作者指出，强大的生成模型能够更好地理解数据，这意味着Masked-LMs的生成能力不足可能也限制了它们的预测能力。 As a famous quote,“what you cannot create, you do not understand.”
  2. 自回归语言模型的不足：
    • 自回归语言模型（AR-LMs）虽然在生成方面具有优势，但在蛋白质序列理解上表现相对较弱。此外，蛋白质是结构性大分子，而非简单的线性符号串，自回归模型只能访问单向上下文，难以捕获氨基酸之间的复杂全局交互，从而影响生成和预测能力。

这些问题显示了对通用蛋白质语言模型的需求，该模型应具备强大的生成和预测能力。

扩散模型的优势

扩散模型在连续数据（如图像）生成上取得了巨大的成功，尤其在高质量图像生成和蛋白质结构建模方面表现突出。扩散模型的优势在于其非自回归的去噪生成方式以及全局感知能力，这些特性特别适合蛋白质结构建模。然而，直接应用高斯扩散模型于蛋白质序列需要额外的连续松弛，无法有效处理离散序列数据。

研究目标

基于上述动机，作者提出了扩散蛋白语言模型（DPLM），旨在通过基于离散扩散概率框架的生成性预训练，打造一个统一且多功能的蛋白质语言模型。DPLM在预训练过程中通过去噪学习蛋白质序列的复杂内在依赖关系，进而用于蛋白质序列生成和下游预测任务。

贡献

作者的主要贡献包括： 1. 提出了一种新的离散扩散蛋白质语言模型，结合了语言模型的可扩展性和扩散模型的强大生成能力，能够完成无条件和条件的蛋白质生成任务。 2. DPLM在生成的新蛋白质序列上具有良好的结构合理性和多样性。 3. 在多种蛋白质预测任务中表现出色，优于现有的主流模型。 4. DPLM支持多种条件生成，包括部分序列、交叉模态以及基于分类器引导的定向生成，展示了其在不同场景下的多功能性。

总结来说，DPLM通过扩散预训练，使模型既能生成具有高度结构合理性的新序列，也能为下游预测任务提供有效的表征。

DPLM: A Versatile Protein LM

f1

3.1 动机（Motivation）

  1. 现有模型的局限性：
    • 目前的蛋白质语言模型未能完全发挥其生成和理解的潜力。掩码语言模型（Masked-LM）虽然在蛋白质序列理解任务中表现优异，但缺乏用于生成任务的有效框架。而自回归语言模型（AR-LM）尽管在生成方面有所优势，但由于其单向上下文限制，无法有效捕获蛋白质序列中的复杂全局依赖。
  2. 扩散模型的优势：
    • 连续扩散模型（如Ho等人，2020；Song等人，2020）在生成连续数据（如图像）方面取得了巨大成功，这激发了作者将扩散模型应用于蛋白质序列生成。扩散模型具有全局感知场和逐步去噪生成的能力，这些特性对于建模蛋白质序列中的复杂相互作用尤为重要。

3.2 使用离散扩散进行蛋白质语言建模（Protein Language Modeling w/ Discrete Diffusion）

  1. 离散扩散建模（Modeling）：
    • 离散扩散过程：
      • 采用类别分布对蛋白质序列进行建模。扩散过程是通过一个马尔可夫过程逐步进行的，在每个步骤中，根据噪声调度参数 β_t 决定是否保持当前氨基酸或者引入噪声： \[ q(x^{(t)}|x^{(t-1)}) = Cat(x^{(t)}; β_t x^{(t-1)} + (1 - β_t) q_{noise}) \]
      • 在该过程中，最初的序列逐渐被加入噪声，直到到达稳定的噪声分布状态。对于正向扩散，从原始序列到噪声的转换可以表示为： \[ q(x^{(t)}|x^{(0)}) = Cat(x^{(t)}; α_t x^{(0)} + (1 - α_t) q_{noise}) \] 其中 \(α_t\) 是正向过程中各步噪声调度的乘积，随着 \(t\) 增加，\(α_t\) 趋于 0，表示序列逐渐被完全随机化。
    • 吸收状态扩散模型：
      • 该模型采用吸收状态作为噪声的稳定分布，使得扩散后的序列可能完全被掩码标记吸收。这与掩码语言模型中的掩码操作类似，即 \(x^{(t)}\) 要么保持为 \(x^{(0)}\)，要么被掩码掉。
  2. 学习过程（Learning）：
    • 学习目标：
      • 学习过程的目标是通过反向去噪过程恢复原始数据分布，优化模型使得生成序列能够最大程度接近输入数据。使用 KL 散度来优化反向模型 \(p_\theta\) 的参数，目标函数可以表示为： \[ J_t = \mathbb{E}_{q(x^{(0)})} \left[ -KL(q(x^{(t-1)}|x^{(t)},x^{(0)}) \parallel p_\theta(x^{(t-1)}|x^{(t)}) \right] \] 通过对离散扩散过程中的各步进行优化，模型学习到如何从含有噪声的输入中重建序列。
    • 与掩码和自回归模型的关系：
      • 离散扩散的学习过程包含了掩码语言模型和自回归语言模型的特例，具体而言，掩码语言模型和自回归模型可以被看作是不同噪声配置下的离散扩散模型，从而统一了两种模型。

3.3 演化级预训练（Evolutionary-scale Pre-training）

  1. 预训练数据集：
    • 采用 UniRef50 数据库进行预训练，该数据库包含约 4500 万条蛋白质序列，总计约 140 亿个氨基酸符号。对于特别长的蛋白质序列，模型会截断至 1024 个符号，以保证计算效率和模型的有效性。
  2. 模型架构：
    • DPLM 的网络架构与 ESM2 相似，采用了 150M、650M 和 3B 三个不同规模的模型。在所有模型中，作者进行了 100K 次更新，150M 参数模型的批大小为 320K，而 650M 和 3B 模型的批大小为 1M。

3.4 生成和表示学习（Generation and Representation Learning）

  1. 生成过程（Generation）：
    • 反向去噪生成：
      • DPLM 的生成过程通过反向迭代去噪实现，从完全掩码的噪声状态开始逐步恢复序列。每次迭代中，模型预测部分掩码符号并更新，从而逐步生成出完整的氨基酸序列。具体而言，在时间步骤 \(t\) 时，模型首先根据 \(p_\theta(·|x^{(t)})\) 生成预测序列 \(x^{\hat{0}}\)，然后通过 \(q(·|x^{(t)}, x^{\hat{0}})\) 生成去噪后的 \(x^{(t-1)}\)。
  2. 表示学习（Representation Learning）：
    • DPLM 在预训练中不仅学会了如何生成序列，还学会了如何从不同噪声水平中去重建输入数据，这使得它在理解蛋白质序列上的表现也非常出色。模型能够对输入的氨基酸序列生成嵌入，这些嵌入可以用于下游的预测任务，如序列或残基级别的分类和回归任务。

3.5 条件生成（Conditioning）

  1. 部分序列条件生成（Conditioning on Partial Sequence）：
    • DPLM 可以在给定部分序列的条件下生成新的序列。例如，它可以根据指定的多肽生成支架以支持功能基序或填充抗体 CDR 环。模型学习从部分观测的序列中进行生成，确保生成的序列保持与指定部分的一致性。
  2. 跨模态条件生成（Adapting DPLM to Condition on Other Modalities）：
    • 对于给定其他模态信息（例如结构信息）的条件生成任务，DPLM 可以结合预训练的模态编码器（例如 GVP-Transformer Encoder），以及新的交叉注意力模块进行微调，生成符合条件的新蛋白质序列。这种跨模态生成方式在应用中非常有用，例如逆折叠任务或小分子配体结合设计等。
  3. 基于分类器引导的可控生成（Plug-and-play Controllable Generation with Discrete Classifier Guidance）：
    • 在某些情况下，直接构建条件模型可能由于数据稀缺而无法实现。因此，作者提出使用离散分类器引导的方法，通过结合预训练分类器来控制生成过程，使得模型能够根据用户偏好生成特定的序列。具体而言，模型在每个生成步骤中结合分类器输出的梯度，对生成的方向进行调整，以确保生成满足用户指定的属性。
  1. ESM2 和 DPLM 的区别：
    • DPLM 在架构和预训练策略上与 ESM2 相似，但在生成能力上有显著提升，无需使用蒙特卡洛方法或吉布斯采样等昂贵的技术。
  2. 与 EvoDiff 的比较：
    • EvoDiff 是另一种基于扩散的蛋白质生成模型。与 EvoDiff 相比，DPLM 不仅展示了更好的表示学习能力，而且具备更广泛的条件生成能力，包括跨模态条件和基于分类器引导的生成，这使得 DPLM 在实际应用中更为灵活和强大。

3.7 小结（Summary）

DPLM 通过基于离散扩散的生成性预训练，成功结合了语言模型的可扩展性和扩散模型的生成能力，不仅可以生成具有高度结构合理性的新蛋白质序列，还能提供有效的表示学习，用于下游的蛋白质预测任务。DPLM 的多功能性使其成为未来蛋白质研究中的重要工具，能够满足多种生成需求，包含无条件和条件生成，同时在蛋白质理解任务中表现优异。

实验（Experiments）

在这一部分，作者通过大量实验评估了 DPLM 的性能，主要从生成和预测任务两方面来验证模型的有效性。具体实验包括蛋白质序列的无条件生成、条件生成以及多个蛋白质表征学习任务。

4.1 无条件生成（Unconditional Generation）

  1. 实验目标：
    • 评估 DPLM 生成的新蛋白质序列的质量，包括序列的多样性和结构合理性。作者使用 AlphaFold2 对生成的序列进行结构预测，以验证其生物学上的合理性。
  2. 实验设置：
    • 在 UniRef50 数据集上进行预训练后，使用训练好的 DPLM 生成全新蛋白质序列。生成过程从完全噪声化的序列开始，模型逐步去噪直到生成最终的蛋白质序列。
  3. 评价指标：
    • 使用 pLDDT（predicted Local Distance Difference Test）评分来衡量生成蛋白质的结构可信度。pLDDT 分数越高，表明预测结构越稳定和可信。
    • 使用 序列多样性（Sequence Diversity）来衡量生成序列的多样性，计算生成序列之间的 Levenshtein 距离。
  4. 实验结果：
    • 生成的蛋白质序列平均 pLDDT 值大于 80，显示出生成序列的结构合理性。
    • 在序列多样性方面，DPLM 生成的序列涵盖了广泛的变异空间，展示了生成序列在不同结构和功能上的多样性。
    • 与 ESM2 和 EvoDiff 相比，DPLM 生成的蛋白质结构更为稳定和多样化。

4.2 部分序列条件生成（Conditioned on Partial Sequence）

  1. 实验目标：
    • 验证 DPLM 在给定部分序列的条件下生成完整序列的能力。此实验旨在测试 DPLM 在蛋白质设计任务中的应用潜力，如生成支持已知功能区域的支架部分。
  2. 实验设置：
    • 实验中对序列进行部分掩码，模型在生成过程中根据给定的部分信息补全剩余部分。使用 AlphaFold2 验证生成的序列是否能够保持给定部分的功能和结构。
  3. 实验结果：
    • DPLM 能够成功地在保持给定功能基序的情况下生成支架部分，显示了其在特定条件下控制生成序列的能力。
    • 实验还显示，DPLM 所生成的序列不仅与给定部分具有良好的兼容性，还能保持整体结构的稳定性。

4.3 跨模态条件生成（Cross-modal Conditional Generation）

  1. 实验目标：
    • 评估 DPLM 在条件生成中的灵活性，尤其是在给定其他模态（如三维结构信息）的情况下生成序列的能力。该实验的重点是测试 DPLM 在复杂应用场景下的可扩展性，如逆折叠任务。
  2. 实验设置：
    • 使用带有三维结构条件的信息进行条件生成。具体地，通过交叉注意力模块集成三维结构信息，指导 DPLM 生成符合目标结构的序列。
  3. 实验结果：
    • DPLM 能够生成符合给定三维结构的序列，其结构精度（measured by RMSD）显著高于无条件生成的基线模型，表明该模型在生成过程中能够很好地利用额外的结构信息。
    • 在多样化生成方面，DPLM 也能在保持目标结构的前提下生成多种不同序列，展示了其在结构约束下生成多样化序列的能力。

4.4 表征学习任务（Representation Learning Tasks）

DPLM 还被用于多种下游的蛋白质预测任务，以验证其生成的序列嵌入在表征学习中的有效性。这些任务包括蛋白质功能预测、亚细胞定位预测和残基级别的特征预测。

f2
  1. 蛋白质功能预测：
    • 实验在蛋白质功能分类任务中，使用 DPLM 生成的序列嵌入进行预测。
    • 结果显示，DPLM 嵌入在各个功能预测基准测试中均表现出色，优于基于 ESM2 和其他主流模型的嵌入。这表明 DPLM 所学习到的序列表征具有较强的生物学意义。
  2. 亚细胞定位预测（Subcellular Localization Prediction）：
    • 使用生成的序列嵌入进行亚细胞定位预测，评估模型在预测蛋白质在细胞内具体位置的能力。
    • 结果表明，DPLM 的表征在这一任务上也显著优于其他模型，特别是在捕捉序列间的复杂相互作用方面具有优势。
  3. 残基级别预测任务（Residue-level Property Prediction）：
    • DPLM 的嵌入还被用于残基级别的预测任务，例如溶解度、二级结构和接触图预测。
    • DPLM 的表征能够很好地捕捉到残基之间的相互作用，结果在多项基准测试中超越了目前最好的模型。

4.5 与现有模型的对比实验（Comparison with Baseline Models）

为了全面评估 DPLM 的性能，作者还进行了与多种现有模型的对比实验，尤其是 ESM2 和 EvoDiff。

  1. 生成任务对比：
    • 与 ESM2 相比，DPLM 在生成任务中的表现更优异，不仅能够生成生物学上可信的蛋白质序列，而且生成的多样性更高。
    • 与 EvoDiff 相比，DPLM 在生成过程中的结构准确性和序列多样性上均有显著提升。
  2. 表征学习任务对比：
    • DPLM 在蛋白质功能预测和亚细胞定位等任务中均超越了基准模型，这主要归功于其扩散预训练过程中学到的深层序列表示。
    • DPLM 所生成的表征更能捕捉到蛋白质中的复杂交互信息，在残基级别的预测任务中表现出色，表明其生成的表征在细粒度的生物学任务中具有较高的实用性。

小结（Summary）

实验结果表明，DPLM 在生成和表征学习任务中均表现优异，展示了其在蛋白质序列生成和理解方面的多功能性和优越性。通过无条件和条件生成实验，验证了 DPLM 在生成生物学上合理的序列方面的能力，尤其是在保持结构和功能的前提下生成新序列的能力。通过表征学习任务，验证了 DPLM 在理解和表征蛋白质序列方面的优势，使其在多种下游任务中表现出色。总体来看，DPLM 结合了扩散模型的生成优势和 Transformer 的表征学习能力，成为通用的、强大的蛋白质语言模型。