[论文精读][计算生物][蛋白质预训练表示]Data-Efficient Protein 3D Geometric Pretraining via Refinement of Diffused Protein Structure Decoy

笔者正在调研市面上的蛋白表示方法，论文方法过于数理的部分会被抽象带过。

Basic Information:

• Title: Data-Efficient Protein 3D Geometric Pretraining via Refinement of Diffused Protein Structure Decoy (通过改进扩散蛋白质结构诱饵来实现数据高效的蛋白质三维几何预训练)
• Authors: Yufei Huang, Lirong Wu, Haitao Lin, Jiangbin Zheng, Ge Wang, Stan Z. Li
  Affiliation: Westlake University, Hangzhou, China (浙江大学计算机科学系, 西湖大学AI研究与创新实验室)
• Keywords: protein structure pretraining, geometric-based, data-efficient, protein-specific pretext task
• URLs: Paper

论文精读

引言

• 蛋白质各种意义上很重要。蛋白序列预训练已经巨大成功，自然语言方法，越来越大的模型。由此想到蛋白结构预训练，且高可信结构预测使得结构预训练变得可能。

• 预训练三要素：数据，模型，训练任务。蛋白特征可以用标量和向量表示。总结：

  1. 序列预训练的大数据，模型和训练任务。预测序列的标量属性。
  2. 结构预训练大幅减少的数据模型和任务。预测标量，有三个问题：样本多样性不足、物理建模不真实和缺乏蛋白质特定预训练任务
  3. 使用三维几何信息进行预训练

• 本文提出数据高效的，基于三维几何的，针对蛋白质的预训练任务，叫做RefineDiff（Refine the Diffused protein Decoy）。模型是方向敏感但旋转不变的transformer架构。仅使用sota 1%的数据就在各种下游任务达到相当的效果。

  文中说是受限于计算能力，但它的数据集来自QA任务，QA任务每个蛋白能生出几百个decoy，，，数据量其实还是很大的。某种意义上也算是数据增强的思路。

相关工作

• 蛋白质结构调整，从能量方法到三维的深度学习
• DDPM去噪声概率扩散模型，启发了RefineDiff。
• 蛋白预训练，ESM家族（序列）；三维残基预训练，将三维结构表示为二维标量输入2D GNN，然后使用2D GNN方法。Uni-Mol三维分子预训练模型表现强大的能力。

蛋白质预训练的统一框架

• 模型，数据，任务，在预训练任务上训练，在下游任务有监督训练，然后测试。

• 蛋白质标量特征应当是旋转不变的

  我的模型对氨基酸的嵌入是旋转不敏感的，但对蛋白整体是旋转敏感的。应该通过数据增强，随机旋转蛋白来强迫模型旋转不敏感。或者其实我的模型的来源任务，识别对象不会发生太多的旋转，但蛋白会。这里可能要考虑到怎么实现全局的旋转不变。

方法

• RefineDiff是一个新颖的、数据高效的、蛋白质特异性的3D几何预训练任务。它包括两个步骤：结构假象的前向扩散过程，然后是蛋白质细化步骤。通过加高斯噪声等方法前向扩散蛋白，然后用模型细化蛋白。这样形成一个自监督。

• 信息瓶颈理论的观点：从信息瓶颈理论的角度看，蛋白质结构假象的前向扩散过程减少了输入和标签之间的相似性，当之后进行结构细化步骤时，模型无法避免通过记忆输入来学习。这可以帮助缓解先前提到的困难的权衡问题，帮助减少学到的表示与输入之间的互信息，增加表示与标签之间的互信息，最终产生更通用和强大的表示。（这段是gpt4概括的）

• 模型结构：强调旋转不变性，用到旋转不变的transformer（这是啥？我能用吗）

实验

预训练

• 蛋白质细化，用DeepAccNet的数据库，限制长度（50-300），最大冗余40%，最差解析2.5埃。 CASP14 Protein Refinement track可以直接作为预训练的测试数据集。

  我的预训练数据没有限制长度，8000原子的超级大蛋白的loss比均值高出两个数量级。老师说这里可以通过结构域人为把蛋白拆开，奈何笔者生物学得不好。暂时用主动学习的思路，把数据丢在一边暂时不处理，但也能观察到这些数据点即使没被纳入训练，它们的loss也在逐epoch下降。

• 细化Alphafold2的输出，只有RefineDiff能提高蛋白得分。说明有效理解了正确的蛋白结构应该长啥样。

  我的预训练效果没法有效衡量。其实应该考虑做一个接触图的副产物，能比较清晰掌握上游预训练的效果。

下游

• EC（酶功能多标签多分类）和GO（基因组学项预测），基本追平SOTA的GearNet，但GearNet用了805K的AFDB，本文只用了8K的数据集（自称DADB）。1%的数据量追平性能。

  这其实好说，毕竟AFDB可以说是PDB的举一反三，DADB可以说是PDB的划重点，两个数据集的信息量差距应当是远没有100倍的。但也表明了这样的三维的蛋白预训练是有效的。

消融实验

• 把RefineDiff换成寻常的Refine，发现寻常Refine的下游迁移能力远不如RefineDiff。
• 有趣的是，使用RefineDiff，即使在上游训练不足，还不能有效完成Refine任务的时候，就已经能有效提升下游任务的表现了。