DPLM初代 ｜统一的序列表征与生成框架
Diffusion Language Models Are Versatile Protein Learners
这篇工作发表在 ICML 2024来自字节跳动研究院和南京大学的合作团队，通讯作者是顾全全老师。
DPLM 将扩散模型引入蛋白质序列建模，成功实现了高质量的序列生成和强大的表征学习能力的统一。不过DPLM仅限于对序列空间的学习和表征，无法直接处理和生成蛋白质的三维结构信息。
而最近刚发表的后续工作 DPLM-2 是一个更令人心潮澎湃的工作。DPLM-2 采用能够同时理解和生成蛋白质的序列与结构。在下游任务和后续的模型性能评价实验中都取得了非凡的成绩。但本文先带领读者一起学习DPLM-初代。
关键词
深度学习｜序列表征｜扩散模型 ｜ 蛋白质语言模型
文献代码
DPLM: https://arxiv.org/pdf/2402.18567 DPLM2: https://arxiv.org/pdf/2410.13782 Github: https://github.com/bytedance/dplm HuggingFace：
150M：https://huggingface.co/airkingbd/dplm_150m 650M：https://huggingface.co/airkingbd/dplm_650m 3B：https://huggingface.co/airkingbd/dplm_3B
目前还无法通过 HuggingFace 的 Transformers 直接调用，需要通过 Github 安装后按照下面的示例代码使用。序列表征需要切换 GitHub 仓库的 Branches。
from byprot.models.lm.dplm import DiffusionProteinLanguageModel
# model_name = "airkingbd/dplm_150m"
model_name = "airkingbd/dplm_650m"
# model_name = "airkingbd/dplm_3B"
dplm = DiffusionProteinLanguageModel.from_pretrained(model_name)
模型速览
模型的具体应用场景？
序列表征 蛋白设计 多种下游任务
模型是否有开源仓库？✔ 模型大小：150M、650M、3B 载入所需要的显存(FP32)：150M(567.62M)/650M(2.43G)/3B(10.59G) (注：暂不考虑 batch_size 的影响，使用 Adam 训练模型的时候大约是载入内存*4)
下游任务表现
1. Thermostability
惜败于SaProt，DPLM(3B): 0.704 (Top 2)，评价标准是Spearman's 相关系数
2. HumanPPI
DPLM(3B): 90.00% (↑ vs SaProt: 86.41%)取得了SOTA，评价标准是Acc
3. Metal Ion Binding
DPLM(3B): 90.00% (↑ vs SaProt: 86.41%)取得了SOTA，评价标准是Acc
4. EC
惜败于SaProt，DPLM(3B): : 0.883 (↓ vs SaProt: 0.884) 评价标准是
5. GO
DPLM均取得了SOTA，评价标准是
6. DeepLoc 蛋白质定位任务
DPLM在二分类任务上取得了SOTA：DLPM 3B: 93.93 (↑ vs SaProt: 93.55)
亚细胞定位任务上惜败于SaProt, DPLM(3B) 85.32 (↓ vs SaProt: 85.57)
两个任务的评价标准是Acc。
目录
背景介绍
双向注意力和因果注意力的区别 DPLM 提出的背景
DPLM 初代:统一的序列表征与生成框架
DPLM的原理和训练过程简述 下游任务和模型表现
讨论总结 扩展阅读
背景介绍
1.1
Bidirectional attention 和 Causal attention 的区别
如果你仔细看过一些基于 Transformer 模型的架构的话你会发现，你会发现无论是 Encoder-only 还是 Decoder-only 他们的基础架构组件是相同的：即都是由多层 Transformer 块串联组成，每个块都包含:注意力计算层、线性层、Layer Normalization，块内部还有残差链接。那问题来了：是什么东西的不同导致模型被叫做某某 coder-only？是因为只用了完整 Transformer 的一半么？
哈哈显然不是一半的问题，Encoder-Only 和 Decoder-Only 的模型核心区别主要就在于注意力屏蔽(Attention Mask) 方式的不同,也就是双向注意力（Bidirectional attention)和因果注意力（Causal attention）的区别。对于因果注意力而言，其Mask矩阵是一个下三角阵，而双向注意力的MASK矩阵是一个秩为一的矩阵(n-1行都可以被第1行线性表示)。
对于一个标准的Scaled dot product attention，即下式：
Mask矩阵的应用是在已经获得了Sofmax得分，即矩阵乘法、缩放、Softmax之后，乘V之前应用的。对于因果注意力而言，注意到下三角阵是满秩的，所以这个应用过MASK的矩阵在乘V得到的结果矩阵必然是满秩的，这意味着信息不会被损失掉。而对于双向注意力来说，由于秩等于一，这个应用过MASK的矩阵在乘V得到的结果矩阵，会损失纬度，此时信息就被压缩掉了。
1.2
DPLM 提出的背景
DPLM 的提出源于蛋白质语言模型面临的两难困境：Encoder-only 模型，如 ESM 系列虽然在序列理解和表征学习方面表现出色，但难以直接用于序列生成；而Decoder-Only 模型虽然善于生成，但表征能力较弱且难以准确捕捉氨基酸之间的长程依赖关系。
DPLM 初代:统一的序列表征与生成框架
2.1
DPLM的原理和训练过程简述
DPLM诞生于一种更大一统的想法，即：一个语言模型可以表示为一个多步扩散模型的形式。
首先考虑一个单步的扩散模型，用ESM做例子。时刻模型会看到输入序列的所有Token，时刻，模型向前传播，一条序列内有15%的Toekn被替换成<MASK>，剩余85%的Token，那么模型的去噪过程就是把15%被替换成<MASK>的Token还原成输入的样子。而这种单步扩散的模式实际上就是ESM作为语言模型在训练时候的样子。
如果你理解了上面这个单步例子，好，那么把这个单步模型变得一般一些。考虑一个多步过程，即：模型还是会看到输入序列的所有Token，时刻，<MASK> 一部分保持住这些<MASK>，在的基础之上继续<MASK>一部分继续保持住这些<MASK>，如此反复，直到时刻所有的Token都被<MASK>掉，模型看到的是一大长串的<MASK>，反向去噪过程就是重新恢复这些<MASK>。
如果这个多步的例子也理解了，那么，这个例子就是一个吸收扩散模型的原理，也就是DPLM的工作原理，如下图所示的那样：
模型在UniRef50数据集上训练(约个Token，使用ESM2的Tokenlizer统计),使用了8*A100集群训练。
不过，作者在训练的时注意到，吸收扩散会导致MASK比例不是固定的值而是一个变量，范围从0%到100%，而传统的Masked语言模型(MLM)Mask比例固定在15%，这样语言模型始终能够获得85%的上下文信息。换句话说，与MLM相比，对于像DPLM这样的吸收离散扩散模型，在极端情况下，几乎所有的token都会被掩码，这意味着模型需要从零开始恢复原始序列中的所有token。而这对于模型会带来严重的学习挑战，在预训练的早期阶段，此时模型尚未获得足够的能力从有限的观察中提取信息特征和相关性。
所以作者采用了使用训练好的ESM2模型改造成扩散模型的样子，之后在根据扩散语言模型的建模目标训练。
2.2
下游任务及模型表现
经过了晦涩难懂的工作原理和训练过程，现在得到了一个序列表征器，那么我们能用它做什么呢？如下图所示：
DPLM能做序列表示任务，诸如序列和残基水平的分类回归等；其次是条件生成任务，这包含三个子任务，含有部分已知序列的条件生成，可以作为其他模态的编码器——作为一个更大型模型的一部分提供序列的特征表示，或者是能够按照特定要求或属性来控制生成过程并且用于判断生成内容是否符合目标属性的分类器。
2.2.1 序列表示任务
序列表示任务上如下表所示：
1. Thermostability
惜败于SaProt，评价标准是Spearman's 相关系数，越高越好:
DPLM(3B): 0.704 (Top 2) SOTA (SaProt): 0.724
2. HumanPPI
DPLM取得了SOTA，评价标准是Acc，越高越好:
DPLM(650M): 86.41% (↑ vs ESM2(650M): 84.78%) SOTA (DLPM 3B): 90.00% (↑ vs SaProt: 86.41%)
3. Metal Ion Binding
DPLM取得了SOTA，评价标准是Acc，越高越好:
DPLM(3B): ACC: 75.94% (↑ vs SaProt: 75.75%) SOTA (DLPM 3B): 75.94%
2.2.2 蛋白质注释任务
1. EC
DPLM(3B): : 0.883 (↓ vs SaProt: 0.884) SOTA (SaProt): 0.884
2. GO
DPLM均取得了SOTA，评价标准是 ，越高越好:
MF子项：SOTA (DLPM 3B): 0.687 (↑ vs SaProt: 0.678) BP子项：SOTA (DLPM 3B): 0.369 (↑ vs SaProt: 0.356) CC子项：SOTA (DLPM 3B):0.463 (↑ vs SaProt: 0.414)
2.2.3 蛋白质定位任务
DPLM在二分类任务上取得了SOTA，亚细胞定位任务上惜败于SaProt,评价标准是Acc，越高越好。
亚细胞定位：SOTA (SaProt): 85.57 (↑ vs DPLM(3B): 85.32) 二分类：SOTA (DLPM 3B): 93.93 (↑ vs SaProt: 93.55)
2.2.4 无条件序列生成任务
1.可折叠性
由 DPLM 生成的蛋白质序列在所有长度上都获得了最高的 pLDDT（图 A）,总体上与PDB库中已知蛋白质结构的统计结果相似。
另外，对 DPLM 生成的序列进行二级结构分析后发现，β-链的比例更高（图D）
此外，尤其是对于超长蛋白质，放大 DPLM 的参数能带来更好的可折叠性(图E)。
2. 序列新颖性
总体来说，DPLM 的 pdbTM 比 EvoDiff 和自然序列要高，如图 B 所示。当蛋白质长度超过 300 时，DPLM 的 pdbTM 得分会降低，而 pLDDT 却保持在 75 以上。这表明 DPLM 有能力在不同长度的序列中采样到结构与 PDB 不相似的序列，随着序列长度的延长，差异会越来越明显。
注：pdb-TM说作者定义的新评价指标，其定义如下：DPLM产生的序列折叠到的结构与PDB中已知结构的进行比较，计算TMScore，取所有TMScore中的最高值。
3. 序列多样性
作者使用inner-TM分数来量化DPLM采样序列的多样性。如图C所示，DPLM的平均inner-TM较低，说明DPLM能合成具有结构多样性的序列。
注：具体来说：对每个采样的候选序列，用ESMFold预测其结构，并计算它与其他序列间的TMscore。计算TMscore的平均值，将这个平均值作为指标。
4. 模型表征能力
如图2F和图2G所示，作者对比了不同架构的语言模型生成的序列折叠后的pLDDT以及不同训练方式得到了的DPLM生成的序列折叠后的pLDDT。
DPLM相比MLM和AR-LM(Decoder-Only)，是最适合蛋白质序列生成的概率框架。DPLM在可折叠性方面优于MLM和AR-LM，验证了作者选择扩散蛋白质语言模型的动机，即：扩散是更合适的蛋白质建模概率框架。
2.2.5 条件生成任务
1. 以部分序列为条件的骨架生成
如下图所示，作者发现DPLM在解决问题数量和平均成功率方面都优于EvoDiff。此外，在DPLM和EvoDiff都能解决的问题中，除了3ixt外，DPLM的成功率均高于EvoDiff。
这表明DPLM在基序支架构建方面表现优异，能在保持motif结构不变的情况下进行骨架生成
2.逆向折叠任务
反向折叠的目标是找到一个能够折叠成给定蛋白质骨架结构的氨基酸序列。
考虑到目前的DPLM只能编码序列信息，所以作者在DPLM的最后一个网络层中加入了结构嵌入，使用GVP-Transformer编码器作为专门的蛋白质骨架结构编码器。然后在CATH 4.2和4.3数据集上评估DPLM的性能。使用氨基酸恢复率（AAR）进行序列评估，而对于结构评估，我们首先使用ESMFold预测生成序列的结构，然后计算pLDDT评分和预测结构与输入结构之间的自一致性TM评分（scTM）。
DPLM 生成的序列能可靠地折叠到给定的骨架结构中。如表2所示，DPLM的性能可以超越或与我们的强基准模型相当，作者认为这源于DPLM在蛋白质序列知识方面的良好学习。当给定结构骨架信息时，DPLM能够利用这一优势，生成结构既合理又与参考结构相似的序列。
2.2.6 二级结构引导的蛋白采样
分类器引导能对生成过程进行灵活控制，无需对每个新条件进行重新训练，尤其适用于标记数据过于有限，无法直接获得条件模型的情况。
图 4 展示了所提出的离散分类器指导有助于引导预先训练好的 DPLM 生成符合从天然蛋白质模板中提取的二级结构注释的样本。这些发现表明，DPLM 具有很强的可编程性，它的全部潜在生成能力都可以通过即插即用的方式实现。
讨论总结
本研究提出了一个基于离散扩散模型的通用型蛋白质语言模型DPLM，成功实现了序列理解和生成能力的统一。
在理论层面，该工作建立了一个完备的概率框架，将传统的Masked语言模型作为特例包含其中，为蛋白质语言建模提供了新的范式。
在技术层面，实现了包括序列片段条件生成、结构信息整合以及基于分类器引导的可控生成等多种条件生成方法，同时提供了支持不同规模(从150M到3B参数)的可扩展模型架构。
扩展阅读
[1] Structured Denoising Diffusion Models in Discrete State-Spaces
文章证明了BERT类的Encoder-Only模型实际上可以看作是一步扩散模型，自回归模型可以看作是特殊的离散扩散模型，让扩散模型与生成式掩码语言模型(MLM)建立了理论联系。 文章地址：https://arxiv.org/abs/2107.03006
[2] Unified Language Model Pre-training for Natural Language Understanding and Generation
NLP领域中的一篇目的相似文章，尽管BERT模型已经显著地提高了大量自然语言理解任务的效果，但是由于它的注意力双向性使得它很难应用于自然语言生成任务。文中为了控制对将要预测的token可见到的上下文，作者使用了不同的self-attention mask来实现一个新的统一预训练语言模型（UniLM）使其既可以应用于语言理解任务，又可以应用于语言生成任务。 文章地址：https://arxiv.org/abs/1905.03197
[3] DiffusionBERT: Improving Generative Masked Language Models with Diffusion Models
模型和MASK策略相似的文章，Table 1对于吸收扩散很直观的能看出区别，是对我在工作原理部分解读的补充。 文章地址：https://arxiv.org/abs/2211.15029
往期合集
干货文章｜线上报告｜蛋白设计｜综述｜Binder｜Co-design｜Benchmark｜AF3｜David Baker｜ESM｜MaSIF｜结构预测｜语言模型｜骨架生成｜逆向折叠｜抗体设计｜多肽设计｜酶设计｜稳定性｜药物设计
关注我们
死磕自己，愉悦大家
专注于*AI蛋白相关的论文解读&学术速运
