ICML 2024 | 扩散语言模型是多才多艺的蛋白质学习者

DRUGAI

今天为大家介绍的是来自字节跳动Quanquan Gu团队的一篇论文。这篇论文介绍了扩散蛋白质语言模型（DPLM），这是一种多才多艺的蛋白质语言模型，展示了对蛋白质序列强大的生成和预测能力。首先，作者在一个生成自监督离散扩散概率框架中，使用进化尺度的蛋白质序列进行可扩展的DPLM预训练，这以一种有原则的方式推广了蛋白质的语言建模。在预训练之后，DPLM展示了无条件生成结构上合理、新颖且多样的蛋白质序列的能力。作者进一步证明了所提出的扩散生成预训练使DPLM对蛋白质有更好的理解，成为一种更优的表示学习者，可以针对各种预测任务进行微调，与ESM2相比具有优势。此外，DPLM可以根据各种需求进行定制，展示了其在多种条件生成方面的强大能力：（1）基于部分肽序列的条件生成，例如高成功率地生成功能基序(functional motif)的支架（scaffold）；（2）将其他模态作为条件，如结构条件生成用于逆折叠；（3）通过插入式分类器引导，将序列生成引导至所需的特性，例如满足指定的二级结构。

尽管当前的蛋白质语言模型已经取得了显著进展，但它们尚未发挥出全部潜力。其中一个基本问题源于广泛使用的预训练目标，即掩码预测与自回归：

(i) 对于掩码预测，掩码语言模型（Masked-LMs，例如ESM系列）在蛋白质预测任务中表现出色，得益于它们的双向感受野。然而，由于缺乏对生成建模的明确定义，掩码语言模型无法进行蛋白质序列生成。作者进一步推测，这甚至可能限制它们的预测能力，因为一个强大的生成模型能够通过学习底层数据分布来创建新样本，同时也会深入理解数据。正如那句名言所说，“你不能创造的东西，你就不理解。”

(ii) 对于自回归，自回归语言模型（AR-LMs，例如ProGen），尽管在生成方面表现良好，但通常在理解序列数据，包括蛋白质方面有所不足。更重要的是，蛋白质是结构性大分子，而不是简单的线性字符串。因此，虽然AR-LMs作为文本归纳偏差是有效的，但它们受限于单向感受野，仅能访问单侧序列上下文。这一限制源于对氨基酸复杂全局相互作用的捕捉不足，从而阻碍了蛋白质语言模型的生成和预测能力。

这突显了对一种通用且多才多艺的蛋白质语言模型的需求，该模型能够结合预测和生成能力。基于上述分析，作者认为，实现这种多功能蛋白质语言模型的关键因素在于：(1) 强大且可扩展的生成建模框架，以最佳方式处理大量蛋白质序列；(2) 双向感受野，以更好地建模逐残基的全局相互作用。而扩散模型在生成高仿真度图像方面大放异彩，还在蛋白质结构建模方面取得了难以置信的成就。它的成功可以归功于扩散模型中具有迭代微调与全局感受野的非自回归去噪生成。

这些因素使得扩散模型成为蛋白质语言建模中具有吸引力的生成基础。然而，直接将传统的高斯扩散应用于蛋白质序列需要额外的连续松弛，这与蛋白质序列的离散性质不符，并且在实践中尚未证明成功。

模型部分

离散扩散（discrete diffusion），即直接在离散状态空间上操作的扩散模型，成为蛋白质序列更适合的概率模型。离散扩散的前向过程定义了一个由转移核控制的马尔可夫过程，其中Cat(x; p)为蛋白质序列x的分类分布。

如Austin等（2021年）所述，离散扩散本质上与AR-LM和Masked-LM相关，而Zheng等（2023a）进一步简化了离散扩散的学习目标，通过他们提出的重新参数化后向转移，将两个分类分布之间的KL散度转化为重新加权的交叉熵：

给定一个训练好的DPLM，它可以通过离散扩散的反向迭代去噪过程合成新的氨基酸序列。形式上，离散扩散从以下分布中采样：

图 1

能够高效采样出真实的蛋白质是必要的，但对于下游应用如治疗开发来说还不够，因为无条件的样本不太可能具备所需的功能特性。这里，作者详细说明了如何通过针对各种需求进行条件设置，使DPLM在实际中更有用，包括最常见的场景，即序列条件设置、跨模态条件设置和即插即用的偏好引导条件设置。

案例一：基于部分序列的条件设置（图1C-1）。包含预先指定多肽的蛋白质生成对应于各种用例，例如为给定的功能基序生成支架、填充抗体CDR环或预先施加专家知识。

案例二：将DPLM适应于其他模态的条件设置（图1C-2）。生成受跨模态约束c约束的蛋白质序列，在实际应用中具有重要价值，例如反向蛋白质折叠，在给定骨架结构的情况下生成序列，或基于小分子配体进行结合物设计。鉴于DPLM主要处理氨基酸标记，在这些情况下，可以通过适配器微调与预训练模态专家编码器Eφ(c)以及新添加的基于交叉注意力的适配器，为DPLM配备跨模态条件设置。

案例三：通过离散分类器引导进行即插即用的可控生成（图1C-3）。

无条件生成性能

图 2

图2显示了DPLM在无条件生成任务中的结果，作者评估了一组长度为[100, 200, ..., 900, 1000]（间隔为100）的性能。DPLM的逆过程采样迭代进行了500步。同时，作者还随机从UniRef50中选取相同长度的天然序列作为参考（标记为UR50）。

(1) 关于可折叠性：DPLM能够生成具有合理预测结构的蛋白质序列。作者使用最先进的单序列结构预测模型ESMFold来检查蛋白质序列的结构合理性或可折叠性，并通过预测的局部距离差异测试（pLDDT）评分来衡量。如果pLDDT > 70，则被认为是高置信度。作者发现，DPLM生成的蛋白质序列在所有长度上都达到了最高的pLDDT评分（图2A）。此外，DPLM生成的序列的二级结构分析显示出更高比例的β折叠（图2D），整体上与PDB中已知蛋白质结构的统计数据相似。而且，作者可以看到，扩展DPLM的规模可以提高可折叠性性能，尤其是对于非常长的蛋白质（图2E）。

(2) 关于新颖性。作者研究了DPLM是否能够生成具有新颖结构的序列，并通过与PDB中的已知结构进行结构相似性比较来进行评估，使用TMScore进行测量。最高的TMScore用于衡量每个序列的新颖性，我们称之为pdb-TM评分。总体而言，DPLM的pdb-TM评分相对高于EvoDiff和天然序列，如图2B所示。有趣的是，当蛋白质长度超过300时，DPLM的pdb-TM评分会下降，但仍保持pLDDT > 75。这表明DPLM具有生成与PDB中各种长度的已知结构不相似的序列的能力，且随着序列长度的增加，这种差异变得越来越明显。

(3) 关于多样性。作者通过inner-TM评分量化DPLM生成序列的多样性。具体而言，对于每个生成的候选序列，作者使用ESMFold预测其结构，并与其余序列计算TM评分。平均TM评分被视为多样性指标。如图2C所示，DPLM的平均内部TM评分相当低，这表明DPLM能够合成结构上多样化的序列。

(4) 关于学习：与掩码语言模型（Masked-LM）和自回归语言模型（AR-LM）相比，离散扩散是最适合蛋白质序列生成的概率框架。如图2F所示，DPLM在可折叠性方面优于掩码语言模型和自回归语言模型，这验证了作者追求扩散蛋白质语言模型的动机，即扩散是蛋白质建模的更适合的概率框架。此外，自回归语言模型在精确控制生成序列长度方面也存在不足，使其在实际应用中不够灵活。如图2G所示，作者发现尽管通过从头开始直接预训练DPLM（DPLM-FS）能够在生成质量上优于ESM2，但这会带来额外的学习挑战和训练开销。因此，作者采用了两阶段训练策略，第一阶段目标为掩码语言建模，随后为扩散目标，解决了这一问题，并获得了接近90的高质量生成pLDDT。

(5) 案例研究。在图2H中，作者展示了DPLM生成的各种长度的蛋白质，从100到1000。随着蛋白质长度的增加，其结构的复杂性也会增加，包含丰富的螺旋和片层结构。作者还发现，DPLM可以生成由串联重复组成的蛋白质，例如β桶或Kelch重复域。

各种蛋白质预测下游任务的性能

作者在多种蛋白质预测任务中评估了DPLM，包括蛋白质功能预测（热稳定性和金属离子结合）、蛋白质定位预测（DeepLoc）、蛋白质注释预测（EC和GO）、蛋白质-蛋白质相互作用预测（HumanPPI），在每个数据集上都进行了全参数的监督微调。作者还包括了来自TAPE的二级结构线性探测。

表 1

DPLM是一个优越的蛋白质序列表示学习模型。正如表1所示，DPLM在所有任务中都优于ESM2。这种性能提升归因于提出的扩散预训练方法，该方法要求DPLM从不同程度的掩码中重建原始序列，包括非常高的噪声水平，而ESM2则使用固定的15%掩码比例。在这种情况下，掩码氨基酸的重建任务变得更加具有挑战性，促使模型从上下文中捕捉到更深层次的依赖关系。此外，作者惊讶地发现，DPLM的表现也接近于SaProt，这是一个基于Foldseek和AlphaFold等折叠模型的结构感知语言模型。这表明DPLM可能从大量的序列数据中隐含地学习到了蛋白质结构。将显性结构信息整合到DPLM中，可能会带来进一步的好处，这值得进一步探索。研究结果证实了作者的初衷，即DPLM通过生成学习过程对蛋白质有了更深入的理解，即通过学习生成蛋白质来更好地理解它们，从而提高了预测性能。

条件生成任务上的性能

图 3

DPLM可以为给定的功能基序生成合理的支架。如图3所示，作者发现DPLM在解决问题的数量和平均成功率方面优于EvoDiff。此外，在DPLM和EvoDiff都能解决的问题上，除了3ixt外，DPLM的成功率均高于EvoDiff。这表明DPLM在基序支架生成方面表现出色，能够在生成支架的过程中保持基序结构。

为了获得更多见解，作者将结构约束的DPLM与在结构设计上达到sota性能的RFDiffusion模型进行了比较。作者发现DPLM在6个问题上表现更好，尤其是在1PRW和5YUI上。利用基序结构帮助DPLM在4个问题上进一步改进了性能，而在其他6个问题上则表现有所下降。这意味着对于某些特定的基序，在序列空间中进行支架生成可能会更好。

表 2

DPLM生成的序列能够自信地折叠成给定的骨架结构。如表2所示，DPLM在氨基酸恢复率（AAR）方面能超越或与强大的基准方法（包括最新的LM-DESIGN方法）持平，更重要的是在结构评估方面表现出色（scTM = 0.85 和 pLDDT > 76）。作者认为，这源于DPLM对蛋白质序列知识的深度学习。在给定结构骨架信息的情况下，DPLM能够利用这一优势生成结构合理且与参考结构相似的序列。

图 4

DPLM具有即插即用的编程能力。图4展示了所提出的离散分类器引导如何帮助预训练的DPLM生成满足提供的二级结构注释的样本，这些注释是从模板天然蛋白质中提取的。这些发现表明，DPLM高度可编程，其生成能力的全部潜力可以以即插即用的方式实现，表明DPLM保留了扩散模型中对离散数据进行可控生成的吸引人特性。这种迅速适应用户在广泛偏好上的不断变化需求的灵活性，在时间和计算至关重要的实际应用中也具有重要意义。

编译 | 黄海涛

审稿 | 曾全晨

参考资料

Wang, X., Zheng, Z., Ye, F., Xue, D., Huang, S., & Gu, Q. (2024). Diffusion Language Models Are Versatile Protein Learners. arXiv preprint arXiv:2402.18567.