[论文解读] ProGen2: Exploring the Boundaries of Protein Language Models

Erik Nijkamp, Jeffrey A. Ruffolo+3 more|arXiv (Cornell University)|Jun 27, 2022

Machine Learning in Bioinformatics被引用 50

一句话总结

ProGen2 将蛋白质语言模型扩展至多达 6.4B 参数，在超过十亿条序列上训练，达到最先进的困惑度，能够生成可行的蛋白质并实现零样本适应性预测。

ABSTRACT

Attention-based models trained on protein sequences have demonstrated incredible success at classification and generation tasks relevant for artificial intelligence-driven protein design. However, we lack a sufficient understanding of how very large-scale models and data play a role in effective protein model development. We introduce a suite of protein language models, named ProGen2, that are scaled up to 6.4B parameters and trained on different sequence datasets drawn from over a billion proteins from genomic, metagenomic, and immune repertoire databases. ProGen2 models show state-of-the-art performance in capturing the distribution of observed evolutionary sequences, generating novel viable sequences, and predicting protein fitness without additional finetuning. As large model sizes and raw numbers of protein sequences continue to become more widely accessible, our results suggest that a growing emphasis needs to be placed on the data distribution provided to a protein sequence model. We release the ProGen2 models and code at https://github.com/salesforce/progen.

研究动机与目标

• 研究极大规模蛋白质语言模型在生成、结构感知生成和适应性预测方面的表现。
• 评估训练数据规模与组成（基因组、微生物组、免疫受体）对模型质量的影响。
• 在不进行任务特定微调的前提下，通过结构预测和功能代理来评估生成质量。
• 在多样化的实验背景和抗体数据集上评估零样本适应性预测能力。

提出的方法

• 基于蛋白质序列的下一个标记预测进行自回归变换器架构训练。
• 模型规模：151M、764M、2.7B 和 6.4B 参数，采用从左到右的因果掩蔽。
• 用于序列位置的旋转位置编码。
• 使用 TPU-v3、JAX 和 pjit 实现数据/模型并行的并行训练。
• 训练数据集：Uniref90、BFD30，以及一个用于专门模型的大型抗体集合（OAS）。
• 通过在保留数据上的困惑度进行评估，对生成序列进行 AlphaFold2 结构预测，进行 Foldseek TMscore 比较，并在多个数据集上进行零样本适应性基准测试。

实验结果

研究问题

• RQ1 增大模型规模如何影响捕捉观测到的蛋白质序列分布的能力（在不同身份阈值下的困惑度）？
• RQ2 在不进行任务特定微调的情况下，是否可以让大型 PLMs 生成新颖且结构上合理的蛋白质和抗体？
• RQ3 预训练模型在多样的蛋白质景观和抗体数据集上是否能提供准确的零样本适应性预测？
• RQ4 在结构聚焦的数据集上进行微调如何影响生成质量及与天然蛋白质的结构相似性？

主要发现

• 更大的 ProGen2 模型产生显著更低的困惑度，表明更好的分布捕获。
• 生成的蛋白质通常折叠成合理的结构并类似天然折叠，尽管存在显著的序列偏差以及偶发的新型折叠。
• 在结构-架构数据集上的微调提高了与天然序列和结构的相似性，采样参数影响多样性与相似性。
• 抗体生成结果显示提示序列与训练分布更对齐，且开发属性（聚集性、溶解性）随采样选择而变化。
• 在狭窄、广阔以及抗体特定的景观中均展示了零样本适应性预测能力，并得到结构和功能评估的支持。