[论文解读] PSP: Million-level Protein Sequence Dataset for Protein Structure Prediction

Sirui Liu, Jun Zhang+17 more|arXiv (Cornell University)|Jun 24, 2022

Protein Structure and Dynamics被引用 9

一句话总结

本文介绍了PSP，这是首个用于蛋白质结构预测的开源百万级别蛋白质序列数据集，包含570,000条真实结构序列和745,000条蒸馏序列。该数据集支持采用基准化训练流程训练最先进模型，在CAMEO竞赛中实现SOTA性能，在51个蛋白质的验证集上达到87.2的lDDT和90.1的TM得分。

ABSTRACT

Proteins are essential component of human life and their structures are important for function and mechanism analysis. Recent work has shown the potential of AI-driven methods for protein structure prediction. However, the development of new models is restricted by the lack of dataset and benchmark training procedure. To the best of our knowledge, the existing open source datasets are far less to satisfy the needs of modern protein sequence-structure related research. To solve this problem, we present the first million-level protein structure prediction dataset with high coverage and diversity, named as PSP. This dataset consists of 570k true structure sequences (10TB) and 745k complementary distillation sequences (15TB). We provide in addition the benchmark training procedure for SOTA protein structure prediction model on this dataset. We validate the utility of this dataset for training by participating CAMEO contest in which our model won the first place. We hope our PSP dataset together with the training benchmark can enable a broader community of AI/biology researchers for AI-driven protein related research.

研究动机与目标

• 为解决当前蛋白质结构预测中AI模型训练所面临的大型、多样化且高质量开源数据集严重缺乏的关键问题。
• 提供一个全面、可扩展且可访问的数据集，以支持像AlphaFold2这样的最先进模型的全规模训练。
• 提供详细、可复现的基准训练流程，包含超参数、训练成本及各阶段精度指标。
• 通过在训练和微调过程中采用新型损失函数和数据采样策略，提升模型性能。
• 通过发布数据和训练协议，使更广泛的社区能够访问先进的蛋白质结构预测工具。

提出的方法

• 构建一个大规模数据集，包含570,000条真实结构序列和745,000条蒸馏序列，覆盖高度多样性和高质量。
• 包含结构预测的完整输入特征：氨基酸序列、MSA、模板以及AlphaFold2的atom37格式的原子级坐标。
• 采用两部分数据组织方式：真实结构数据集和蒸馏数据集，每部分均划分为256个压缩块，以实现高效存储和访问。
• 开发一个基准训练流水线，包含详细的超参数、训练成本估算以及各训练阶段的预期精度。
• 在初始训练阶段引入一种新型结构相关损失函数，并优化数据采样策略，以增强模型泛化能力。
• 在微调阶段应用稳定化的违反损失（stabilized violation loss），以提升模型置信度和结构准确性。

实验结果

研究问题

• RQ1 一个百万级别的蛋白质序列数据集能否显著提升最先进蛋白质结构预测模型的训练效果和性能？
• RQ2 一个包含详细超参数和阶段精度指标的全面开源训练基准，在多大程度上能提升AI驱动蛋白质研究的可复现性和可及性？
• RQ3 高质量蒸馏序列和多样化真实结构的引入，如何影响模型的泛化能力和预测精度？
• RQ4 损失函数和数据采样策略的哪些改进能带来结构预测模型性能的提升？
• RQ5 一个面向社区开放的数据集和训练流程，能否在CAMEO等真实世界基准中实现具有竞争力的性能？

主要发现

• PSP数据集包含570,000条真实结构序列（10TB）和745,000条蒸馏序列（15TB），是同类中规模最大的开源数据集。
• 基准训练流程支持最先进模型的全规模训练，包含详细的超参数、训练成本估算以及各阶段的预期精度。
• 在PSP上训练的模型在CAMEO一个月期验证集（51个序列）上达到90.1的TM得分和87.2的lDDT，优于AlphaFold2和RoseTTAFold。
• 使用新型结构相关损失函数和改进的数据采样策略，在初始训练阶段带来了可测量的性能提升。
• 在微调阶段采用稳定化的违反损失，显著提升了模型置信度并减少了结构错误。
• 数据集被组织为256个压缩块，实现高效存储与访问，且与标准蛋白质结构预测流水线完全兼容。