[论文解读] RITA: a Study on Scaling Up Generative Protein Sequence Models

Daniel Hesslow, Niccoló Zanichelli+3 more|arXiv (Cornell University)|May 11, 2022

Machine Learning in Bioinformatics被引用 59

一句话总结

RITA 引入自回归蛋白质序列模型，参数规模高达 1.2B，使用超过 280M 的 UniRef-100 序列进行训练，并系统性研究扩展如何影响在下一个氨基酸预测、适应性预测以及酶功能预测等任务的性能。

ABSTRACT

In this work we introduce RITA: a suite of autoregressive generative models for protein sequences, with up to 1.2 billion parameters, trained on over 280 million protein sequences belonging to the UniRef-100 database. Such generative models hold the promise of greatly accelerating protein design. We conduct the first systematic study of how capabilities evolve with model size for autoregressive transformers in the protein domain: we evaluate RITA models in next amino acid prediction, zero-shot fitness, and enzyme function prediction, showing benefits from increased scale. We release the RITA models openly, to the benefit of the research community.

研究动机与目标

• 推动设计大规模自回归蛋白质序列模型，以加速蛋白质设计。
• 通过分析模型规模增加时的性能来建立蛋白质序列建模的扩展规律。
• 在下游任务上评估 RITA：下一个氨基酸预测、零-shot fitness（ProteinGym）以及酶功能预测。
• 向社区发布模型，以实现开放研究与基准测试。

提出的方法

• 使用预训练数据中的 150B 氨基酸，训练四个尺寸（Small 到 XLarge）的小到特大的解码器仅 Transformer 模型，参数规模高达 1.2B。
• 在对 AliBi 进行消融后使用 RoPE 循环位置嵌入；为降低语言模型损失而选择 RoPE。
• 主要在 UniRef-100 上进行预训练，辅以 Metaclust 和 MGnify 数据，无需预聚类。
• 在 100 多张 GPU 上使用 Megatron-Deepspeed 对 150B 氨基酸进行训练，使用 Adam 优化器，批量大小 512，上下文长度 1024。
• 在 UniRef-100、MGnify 以及 Metaclust 的保留集蛋白家族上评估困惑度；与 ProtGPT2 和 ProtXLNet 进行比较。
• 通过以下方式评估下游任务：(i) 通过 ProteinGym 替换进行突变效应预测；(ii) 使用 SwissProt 注释进行酶功能预测；(iii) 通过提示微调来生成目标蛋白家族。

实验结果

研究问题

• RQ1 自回归 Transformer 中，蛋白质序列模型的能力如何随模型规模的增加而扩展？
• RQ2 与基线相比，较大的 RITA 模型是否在下一个氨基酸预测、适应性预测和酶功能预测方面有所提升？
• RQ3 提示微调是否能进一步提升对未包含在训练集中的蛋白家族的生成质量？
• RQ4 在计算资源约束下，蛋白质序列建模的扩展规律特征（如指数）是什么？

主要发现

• 在所有测试数据集上，困惑度随模型规模提高而改善；RITA-XL 在 UniRef-100、Metaclust、MGnify 以及 Pfam 保留集上达到最佳困惑度。
• 在突变效应（ProteinGym）预测中，RITA 模型的性能随规模增大而提升，对于 Large 与 XLarge 变体甚至可以超过 ESM-1v。
• 在酶功能预测中，模型越大 top-k 准确率越高，XLarge 取得最佳结果（研究中给出从 top-1 到 top-10 的数值）。
• 与基础模型相比，提示微调显著降低了一个未包含的蛋白家族 PF03272 的困惑度，显示了可控生成。
• 在所有下游任务中，RITA 模型接近或超过任务特定基线，并呈现与计算驱动的 NLP 扩展规律相一致的扩展行为，尽管由于词汇表规模较小和训练动态存在显著偏差。