[论文解读] Data-Efficient Protein 3D Geometric Pretraining via Refinement of Diffused Protein Structure Decoy

Yufei Huang, Lirong Wu+4 more|arXiv (Cornell University)|Feb 5, 2023

Protein Structure and Dynamics被引用 5

一句话总结

本文提出 RefineDiff，一种数据高效的3D几何掩蔽任务，用于蛋白质结构预训练，通过旋转和平移不变的Transformer模型对扩散的蛋白质结构伪像进行优化。尽管仅使用了SOTA模型所用数据的不到1%，该方法在下游任务中仍取得了具有竞争力的性能，包括在GO-CC任务上达到SOTA，在EC和GO-BP任务上接近SOTA，证明了3D几何预训练的有效性。

ABSTRACT

Learning meaningful protein representation is important for a variety of biological downstream tasks such as structure-based drug design. Having witnessed the success of protein sequence pretraining, pretraining for structural data which is more informative has become a promising research topic. However, there are three major challenges facing protein structure pretraining: insufficient sample diversity, physically unrealistic modeling, and the lack of protein-specific pretext tasks. To try to address these challenges, we present the 3D Geometric Pretraining. In this paper, we propose a unified framework for protein pretraining and a 3D geometric-based, data-efficient, and protein-specific pretext task: RefineDiff (Refine the Diffused Protein Structure Decoy). After pretraining our geometric-aware model with this task on limited data(less than 1% of SOTA models), we obtained informative protein representations that can achieve comparable performance for various downstream tasks.

研究动机与目标

• 为解决2D蛋白质结构预训练的局限性，包括数据同质性、物理上不现实的建模以及掩蔽任务有限的问题。
• 开发一种3D几何、数据高效且针对蛋白质的掩蔽任务，用于预训练，以利用结构生物学的洞见。
• 证明在有限数据上使用3D几何掩蔽任务进行预训练，可获得与使用大规模数据集训练的蛋白质表征相媲美甚至更优的表征。
• 建立一个统一的蛋白质表征学习框架，凸显3D几何预训练的潜力。

提出的方法

• 提出RefineDiff，一种新颖的掩蔽任务，训练模型将扩散的蛋白质结构伪像优化为更精确的构象。
• 采用前向扩散过程将蛋白质结构扰动为伪像，使模型能够从更广阔的势能景观中学习。
• 采用旋转和平移不变的Transformer架构（受AlphaFold启发），确保模型具备几何感知能力。
• 在仅包含8,000个蛋白质靶标的较小数据集（DADB）上预训练模型，数据量不足SOTA模型的1%。
• 采用去噪目标，即模型从被破坏（扩散）的版本中预测原始、优化后的结构。
• 整合归纳偏置，如三角注意力和几何约束，以提升物理合理性。

实验结果

研究问题

• RQ1 3D几何掩蔽任务是否能比2D几何或基于序列的方法更有效地提升蛋白质表征学习？
• RQ2 在小而多样的数据集上使用3D几何掩蔽任务进行预训练，是否能获得与使用大规模数据集训练的表征相媲美甚至更优的表征？
• RQ3 RefineDiff与直接蛋白质结构优化作为预训练目标相比，在下游任务性能上表现如何？
• RQ4 扩散过程在多大程度上通过逃离局部极小值来增强结构优化？

主要发现

• RefineDiff预训练模型在GO-CC任务上达到SOTA性能，F_max为0.528，优于所有基线模型，包括在更大数据集上训练的模型。
• 该模型在EC和GO-BP任务上分别取得第二名的性能，F_max分别为0.865和0.455，尽管仅使用了8,000个蛋白质靶标（不足SOTA模型数据的1%）。
• 消融研究显示，即使模型的结构优化能力初始较弱，使用RefineDiff进行预训练仍优于直接在蛋白质结构优化上进行预训练。
• 使用RefineDiff训练的模型相比初始伪像结构提升了0.12 GDT，而直接PSR预训练仅提升了0.00 GDT，表明基于扩散的方法具有明显优势。
• 在预训练过程中使用早期检查点仅带来0.01 GDT的提升，表明必须完成完整训练才能释放RefineDiff的全部优势。
• 该模型在所有下游任务上均优于从零开始训练的基线模型，证明即使数据有限，预训练仍具有显著价值。