机器之心发布
机器之心编辑部近日,百图生科与清华大学联合提出了一种名为 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) 已经成为主流范式,通过把数据标签映射到整个文本空间,结合指令微调来生成各式各样的任务的答案,但 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,即 [MASK])。这样的设计确保了 xTrimoPGLM 能生成精确的残基级和序列级表示。当使用 [MASK] 标识符时,xTrimoPGLM 的功能类似于 BERT。相反,当使用 [sMASK] 或 [gMASK] 时,xTrimoPGLM 的行为类似于 PrefixLM 或 GPT。总的来说,xTrimoPGLM-100B 的预训练阶段可以分为两个阶段。首先,利用 MLM 进行预训练以增强其表示能力,主要目标是快速减少损失水平。第二阶段,使用结合 MLM 和 GLM 损失的统一目标进行训练,以提升理解和生成能力。 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 目标函数;
- CDR3 短序列掩码(CSM):对 CDR3 部分区域进行掩码重设计。
- CDR3 全序列掩码(CWM):对 CDR3 全部区域进行掩码重设计。
- CDR3 随机突变(CRM):对 CDR3 区域内特定位点的随机 3-6 个位置进行突变。
- CDR3 随机检索(CRR):使用来自 SARS-CoV-2 野生型库中其他抗体的序列随机替换 CDR3 区域。
© THE END
转载请联系本公众号获得授权