他山之石，可以攻玉——深度解读「千亿级蛋白质语言模型 xTrimoPGLM」
近日，百图生科与清华大学联合提出了一种名为xTrimo Protein General Language Model (xTrimoPGLM)的模型，参数量高达千亿（100B）。相关成果于2023年7月7日在biorxiv上发布。
自然语言处理（NLP）领域中的预训练模型相关进展已经被成功地用于学习蛋白质序列中隐藏的生物信息。现在大多数的蛋白质预训练模型都受限于自动编码或自回归预训练目标，这使得它们难以同时处理蛋白质理解（例如，蛋白质结构预测）和生成任务（例如，药物设计）。
这篇论文提出统一的蛋白质语言模型，xTrimoPGLM，通过一个创新的预训练框架来同时处理这两种类型的任务。xTrimoPGLM主要技术贡献是探索了这两种类型目标之间的兼容性以及共同优化的可能性，并基于此训练了一个前所未有的1000亿参数规模的蛋白质语言模型，并消耗了1万亿Tokens，模型FLOPs达到6.2e+23，达到和175B参数的GPT-3一个量级。
在理解任务上，xTrimoPGLM在多种蛋白质理解任务（15项任务中的13项任务）中显著优于其他先进基线。在生成任务上，xTrimoPGLM能够生成与自然蛋白质结构类似的新蛋白质序列。
此外，文章基于相同的框架额外训练了一个12亿参数的抗体模型（xTrimoPGLM-Ab），其在预测抗体自然性和结构方面取得了市面上最好的效果，并且显示出比AlphaFold2更快的推理速度（数十倍到数千倍）。综合来看，这些结果充分展示了xTrimoPGLM在理解和生成蛋白质序列方面的强大能力和广阔的应用前景。
不同类型的蛋白质相关任务需要蛋白质语言模型（Protein Language Model，PLM）提供各异的输出。具体的，蛋白质理解任务，如二级结构预测等，需要PLM提供精确的氨基酸和序列级别的表示；而蛋白质设计任务，如抗体或酶的设计，依赖于PLM的生成能力。然而，当前的PLM因为其单一的预训练框架的限制，大多只能处理一种类型的任务。
事实上，蛋白质的理解和生成都反映了蛋白质数据的分布信息，Meta之前使用ESM（基于Masked Language Model的蛋白质大模型）做生成的工作也证实了这一点，指出蛋白质理解预训练模型可以通过一些采样策略进行蛋白质设计。这进一步支撑了这两种看似不同的任务的统一性，如果能够使用同一个训练框架去处理这两种任务，将会进一步增强模型对蛋白质数据的拟合能力。
虽然NLP领域生成式模型（例如UL2R, GPT) 已经成为主流范式，通过把数据标签映射到整个文本空间，结合指令微调来生成各式各样的任务的答案，但PLM还无法实现这一点。实际上，蛋白质的应用仍然依赖于将表示与下游任务特定标签（如结构预测的3D坐标）之间的桥接，这在很大程度上依赖于BERT样式的训练来处理蛋白质理解任务。因此，需要同时进行这两种训练目标。
本文介绍的xTrimo蛋白质通用语言模型（xTrimoPGLM）预训练框架，巧妙地统一了两种类型的预训练任务，使模型能同时处理多种与蛋白质相关的任务。研究团队通过全面的实验评估了xTrimoPGLM框架的有效性。在蛋白质理解任务的情况下，xTrimoPGLM-100B在多种评估中表现出色，涵盖了蛋白质结构、功能、交互和可开发性等领域的15项任务。
"Scaling Law"是衡量大型语言模型的重要原则，模型的性能应随着模型参数大小、数据量、计算量按比例指数增加而线性增长。实际上，研究团队在下游任务上的实验结果验证了这一定律，证明了大型模型在处理复杂任务时的必要性。
如下图所示，性能改善与预训练计算量增加之间的关系。以Meta的ESM-2为参考，随着蛋白质语言模型（PLM）的计算量呈指数增长，蛋白质的下游性能仍然会线性增长（每个大类有3-4个任务，数值表示这些任务的平均值）。
在蛋白质生成任务中，xTrimoPGLM-100B展示了生成不同长度和序列的新蛋白质序列的能力，这是通过调整生成超参数实现的。值得注意的是，当与自然蛋白质进行比较时，xTrimoPGLM-100B展示了生成结构相似但序列相异的新蛋白质序列的能力。这再次验证了大型模型对于复杂任务的重要性，进一步证实了遵循"Scaling Law"进行模型设计的决定是正确的。
同时，研究团队还开发了一种当前更具实用意义的具有12亿参数的抗体PLM，即xTrimoPGLM-Ab-1B。这种模型在OAS抗体数据库上进行fine-tuning，处理了超过1万亿个token。它在抗体的自然性和结构预测任务上达到了目前最优秀的性能。由于不依赖于外部库的检索和多序列对齐（Multiple Sequence Alignment)，所以结构预测在速度上比Alphafold2模型提升成百上千倍，这对于基于抗体药物发现AI制药公司至关重要。
原始的GLM模型利用两种不同的预训练目标来提升其生成能力：1）跨度生成（Span Generation，简称sMask），用于恢复句子中的短空白；2）长文本生成（Long-text Generation，简称gMASK），用于在提供前缀上下文的基础上生成随机长度的序列。为了进一步提升xTrimoPGLM的理解能力，团队在prefix区域引入了被用作理解目标的Masked Language Model（MLM，即
当使用
NLP领域大量探索了统一的预训练模式，但大多还是采样了同样的训练模式（自回归或自编码）。为了满足统一的蛋白质预训练模型的需求，需要将BERT样式的目标引入到预训练语言模型中，以增强模型的表示能力，同时也需要引入GPT样式的目标，以确保模型的生成能力。在最开始研究团队使用Probing 策略探索Contact Map Prediction的任务时，发现仅仅依靠基于下一个词预测的生成式语言模型，效果会有大幅度的下降。
兼容性实验：在实证分析中，研究团队在xTrimoPGLM-150m模型上探究了同时优化两个不同目标的可行性。结果如下图(a)(b)所示，即使两种预训练目标看似冲突，MLM损失和GLM损失也可以同时优化，反之亦然，即统一的训练可以很快的适配到MLM或者GLM上，并不会影响两者的收敛情况。
研究团队还探究了MLM与GLM两种目标是否能相互加速收敛，结果如图(c)(d)。
MLM-adapted GLM：接着MLM预训练后的模型，继续训练GLM目标函数；
GLM-adapted MLM：接着GLM预训练后的模型，继续训练MLM目标函数；
总的说来，基于预训练后的模型中继续训练的模型，与从头开始训练的模型相比，其收敛速度明显加快。这些实验观察到：在蛋白质数据分布并不依赖于特定的训练模式，从而缩小了自编码PLMs（如ESM）和自回归PLMs（如ProGen2）之间的差距，为100B训练流程提供了支持。
训练稳定性是成功训练100B规模大型语言模型的决定性因素。xTrimoPGLM从GLM-130B的实现中借鉴了一些想法并解决了许多不稳定训练的问题。然而，xTrimoPGLM-100B在从训练的第一阶段过渡到第二阶段时仍然会遇到灾难性的训练崩溃（小规模的模型（10B规模）中并未观察到），即使一开始只将1%的GLM损失加入预训练也可能触发这些崩溃。
下图可以看到，如果直接给GLM分配一个比例 ，在训练初期，grad norm都会出现spike（橙，蓝，绿线）。
为了缓解这个问题，研究团队提出了一种平滑过渡策略。主要分为两个阶段。在第一阶段，主要目标是逐步提高GLM损失的比例，达到预期的数量。具体来说，给定一个期望的GLM损失比例R，按照线性增长，以K步从0增加到R。在这个阶段应该将学习率保持在极低的水平。完成过渡后，学习率可以按照预定义的脚本在几百个steps内逐渐回升至原来的水平 （紫线）。实际上，最后的xTrimoPGLM-100B训练运行只在过渡阶段经历了损失分歧情况，但是由于硬件故障多次失败，导致经常性的更换节点和重启。
为了能够尽可能的映射整个蛋白世界， xTrimoPGLM-100B的预训练模型的训练数据集整合自两个广泛的数据源：Uniref90和ColAbFoldDB。
结合这两个数据源，预训练模型数据集充分利用了这两个数据源的优势，既有广泛的生物分类覆盖，又有多样的环境生态位蛋白质序列，全面而详尽地映射了生物世界中的蛋白质资源。
xTrimoPGLM-100B模型的训练过程复杂且耗费资源，团队耗费了160天的时间，开发团队使用了96台DGX-A100 GPU服务器（每台服务器拥有8×40G的GPU）以混合精度（FP16）进行训练，消耗了1万亿的tokens，由于大多数现有的大型语言模型在训练上存在严重不足，所以开发团队现在仍然在继续训练xTrimoPGLM-100B模型，以处理尽可能多的tokens。
模型采用了3D并行策略，基于DeepSpeed进行了4路张量并行、8路流水线并行，并且采用了Zero Stage 1进行训练。模型具有72个Transformer层，80个head，以及10,240维和31,744个FNN的维度。使用了DeepNorm初始化Post-LN，并且采用了Embedding Layer Gradient Shrink (EGS)稳定训练，以及使用了2D ROPE位置编码技术。为了提升训练效率，每个样本由多个蛋白拼接再一起，使用<eos>区分开来，包含固定的2,048的序列长度。最终单卡的TFLOPs在120-135之间，68 examples/sec, 如果是80G的A100，经过减少重算可到92 examples/sec。下表展示了大部分超参数的配置。
团队也对比了当前比较热门的预训练模型的FLOPs，目前生物领域不同模型的结构也存在比较大的区别， 所以在FLOPs的计算上，团队考虑尽可能详细统计，包括query、 key、value的转换，Attention Matrix的计算，注意力之后线性变换，MLP中的变换，以及最后到vocab维度的映射，以及语言模型头（如果存在的话）中的线性转换等, 可以看到，xTrimoPGLM-100B高出其他模型一到两个量级。
为了全面评估xTrimoPGLM-100B模型，研究团队对4个领域内，15个下游的蛋白质相关任务进行了基准测试。测试结果表明，xTrimoPGLM-100B在蛋白质结构、蛋白质可开发能力、蛋白质相互作用和蛋白质功能等四个主要类别的任务中，都表现出了显著的优势。在这些任务中，xTrimoPGLM-100B模型与微调技术的组合取得了优异的成绩，大部分超越了当前最先进的方法，从而推动这个领域的进步。
需要强调的是，下表对比主要从任务的角度进行，而不是一个完全公平的对比，因为xTrimoPGLM-100B在取得这些结果时，采用了这个领域大模型之前都不太关注微调技术。这些结果的大部分来自对论文直接引用，并使用相同的数据划分策略，有一些没有benchmark的任务，研究团队使用了ESM-15B + finetuning的策略作为benchmark， 实际上，在所有任务中，研究团队也使用过这种策略，发现ESM2-15B/3B + finetuning在不少任务可以直接达到SOTA，但是，目前大部分的蛋白质大语言模型很少关注微调技术，更多的是把PLMs作为特征提取器使用。
另一方面，为了表明xTrimoPGLM大模型的有效性，研究团队也给出来和 ESM2-15B 以及 150M 的在相同训练设置下的对比实验，使用相对较小的ESM2-150M模型作为指标，主要来理解各种下游蛋白质相关任务的难度程度。这些任务包括基于feature-based 的Probing和 联合大模型参数的Finetuning，xTrimoPGLM-100B在大多数蛋白质相关任务中仍然展现出了优势。
总的来说，xTrimoPGLM-100B在15项任务中的12项上表现得比ESM2-15B更好。研究结果也揭示了一个规律：尽管其它的方法（比如Ankh）在试图寻找一种途径，即在不依赖大规模语言模型的情况下，构建高效、低成本、有知识指导的蛋白质语言模型，但在模型的性能还是与模型规模密切相关，这表明，扩大模型规模可能是一个简单而有效的方法，能够在通用的蛋白质相关任务中提升模型的表现。这为未来对蛋白质预训练模型进一步的研究提供了指导。
为了进一步确定xTrimoPGLM框架的通用性，团队把该框架应用在抗体蛋白预训练上。考虑到训练资源的限制和抗体数据的多样性不足（大部分长度相似且有相似的框架区域），团队没有直接在xTrimoPGLM-100B上进行精调，而是构建了一个12亿参数模型xTrimoPGLM-Ab-1B，在包含10亿抗体序列的OAS数据集对模型进行训练。
考虑到CDRs是抗体最重要的部分，团队对40%的样本进行完整的CDRs掩码处理，另外40%的样本随机掩码处理，而剩下的20%则使用MLM目标。由于
对于antibody-based的药物设计，有两项必不可少的任务就是序列自然度 (Naturalness) 以及抗体结构预测 (Antibody Structure Prediction), 下面分别介绍。
研究团队使用了百图生科湿实验室获得的蛋白质表达实验数据集来评估各种模型的性能。具体来说，任何产生的纯化蛋白质少于10 mg/L的样本都被归类为未表达，而产生超过10 mg/L的样本被认为是成功合成的。第一个数据集（数据集1）包括601个抗体序列，来自在CHO和HEK293细胞上进行的湿实验。其中，成功表达的有516个。第二个数据集（数据集2）包含了98个针对特定抗原的人类抗体序列，其中90个成功表达。评估采用zero-shot方式评估，不针对标签微调，仅通过计算序列困惑度（PPL）和伪困惑度（PPPL）给序列打分。
结果显示，在这两个数据集中，xTrimoPGLM-Ab-1B均优于其他基准模型。而且，进一步对xTrimoPGLM-Ab-1B进行精调，分别得到了xTrimoPGLM-Ab-1B-GLM和xTrimoPGLM-Ab-1B-MLM两个模型。结果显示，这两个模型在数据集2上的AUC得分均有0.02的提升。
作者们还进行了消融研究，以证明随机区域掩码和CDR区域掩码的重要性。实验证明，同时使用这两种目标在数据集1和数据集2上的表现均优于仅使用其中一种任务的模型，这突显了组合使用这两类目标的重要性。
这个任务目标是根据抗体的序列来预测其结构，实验涵盖了单链结构预测和复杂结构预测，即VH-VL复合物。
单链结构预测的数据集源自2022年4月13日之前的RCSB Protein Data Bank (PDB)。该数据集包含19k个抗体链（VL或VH）。通过过滤，最终获得了约7.5k个独特的序列。另一个数据集，VH-VL复合物，包含了大约4.7k个来自PDB的抗体。评估标准为根均方偏差（RMSD）和TM-score。复杂结构预测还包括DockQ评估。
相比当前流行的结构预测模型（如ESMFold, AlphaFold2），xTrimoPGLM-AbFold做了以下改变：1) 去除了MSA和模板搜索模块；2) 将下游evoformer模块的数量从48减少到1。
对于单链结构预测任务，研究团队对比了Alphafold2以及四个基于PLM的模型：OmegaFold、ESMFold、IgFold和xTrimoAbFold。
结果如表所示，xTrimoPGLM-AbFold在所有抗体结构预测相关的指标上都显著优于其他模型，进一步说明，在预训练模型对数据分布拟合的足够好时，只需微调一个额外的Evoformer模块以及不依赖MSA和模板的情况下，就能成为领先的抗体结构预测模型。
对于VH-VL复杂结构的预测，研究团队比较了ZDock、ClusPro、EquiDock、HDOCK以及AlphaFold-Multimer。
上表也展示了不同模型在VH-VL复合物性能上的表现。AlphaFold-Multimer使用了MSA和模板信息，效果优于大多数结构预测算法。而xTrimoPGLM-AbFold不使用任何MSA或模板信息，与AlphaFold-Multimer的表现相当，这说明xTrimoPGLM-Ab-1B已经学习了足够丰富的抗体信息。更重要的是，其推理速度比AlphaFold-Multimer快了6300倍，而且比使用了MSA加速搜索策略的AlphaFold-Multimer快了103倍。在AI制药引擎中，往往需要对生成的候选序列快速进行结构预测，以便算出reward，进行下一轮的迭代，速度的提升决定了引擎的效率。
此外，团队将Evoformer模块的数量增加到16个时，xTrimoPGLM-AbFold在所有指标上都达到了最佳性能，同时速度还比原来的AlphaFold-Multimer快2400倍，比加速MSA搜索版的AlphaFold-Multimer快40倍。值得注意的是，当Evoformer模块的数量从1增加到16时，效果只有少量的提升，这表明预训练模型已经学习到了足够的序列信息，可以精确地预测原子位置。
在探索xTrimoPGLM生成自然功能序列的能力上，研究团队生成了数千个序列并预测其对应的三维折叠结构。研究团队发现该模型能够生成重要的二级结构，包括alpha螺旋和beta折叠，这些都是更复杂三级结构的基础。此外，模型生成的序列与自然序列相似性低，这为之后的药物合成提供了更多的选择。
但是，如何生成高质量序列还有着巨大的挑战。首先，当处理超过200个氨基酸的序列时，模型往往生成大量的循环，而不是自然蛋白质般的结构。其次，模型在生成过程中经常会产生重复的问题，这可能源于模型倾向于选择在局部最大化输出概率的选项。对此团队尝试使用了n-gram penalty来减少生成重复序列的可能性，但研究团队发现许多示例都展示出低复杂度的序列（例如，局部重复），预测的结构中包含长循环无序区域，推测是n-gram惩罚可能阻碍了模型生成正确序列的能力（下图第一行）。在去除n-gram penalty后，模型能够生成正常的结构（下图第二行）。
团队近一步展示了xTrimoPGLM-Ab-1B模型的生成能力，选择了一个能与SARS-CoV-2-WT链接的重链抗体序列，并采用四种不同的策略重新设计了该序列的CDR3区域。因为CDR3区域在抗体或T细胞受体的结构中起着关键作用，并且具有显著的变异性，对抗原识别的特异性起着重要的作用。以下是定义的四种策略：
CDR3短序列掩码（CSM）：对CDR3部分区域进行掩码重设计。
CDR3全序列掩码（CWM）：对CDR3全部区域进行掩码重设计。
CDR3随机突变（CRM）：对CDR3区域内特定位点的随机3-6个位置进行突变。
CDR3随机检索（CRR）：使用来自SARS-CoV-2野生型库中其他抗体的序列随机替换CDR3区域。
团队通过xTrimoPGLM-Ab-1B生成了一组6,000个抗体。研究团队随机选择了六种抗体，并使用xTrimoPGLM-AbFold作为结构预测模型。CSM和CWM策略能够生成不同长度的序列，而不必进行突变或删除。相比之下，由两个并行基线CRM和CRR生成的序列显示出相当大的无序性，无论是否存在少量突变或完全替换整个CDR3片段。研究团队的分析进一步发现，编辑距离与生成抗体CDR3区域的结构之间存在关系。特别是，随着编辑距离的增加，CDR3区域的结构倾向于退化，即使是大的生成模型目前也仍然面临限制。
通过借鉴NLP/CV领域的想法，生物领域的预训练模型近两年雨后春笋般地冒出来，随着模型计算力的提升和生物数据增长，我们期待更多未知的、惊人的发现出现在这个领域中。
尽管仍然存在不少需要继续探索的地方，千亿模型的诞生不仅标志着最前沿的AI技术和生物学技术的融合，还意味着一个充满无限可能的未来已经开启。我们期待，这一重量级的模型引领制药领域步入一个新的黄金时代，为人类健康和科学事业开创更加光明的未来。
论文链接：https://www.biorxiv.org/content/10.1101/2023.07.05.547496v3