使用Vabs-Net进行多层次蛋白质预训练

DRUGAI

今天为大家介绍的是来自Shuqi Lu团队的一篇论文。近年来，在各种下游任务中基于3D结构的预训练蛋白质模型的发展激增，这代表了预训练蛋白质语言模型的重大进步。然而，大多数现有的基于结构的预训练模型主要关注残留物水平，即α碳原子，而忽略了如侧链原子等其他原子。作者认为，在残基和原子水平上对蛋白质进行建模是很重要的，因为侧链原子对许多诸如分子对接等下游任务也至关重要。然而，作者发现在预训练中天真地结合残基和原子信息通常会失败。一个关键原因是输入中包含原子结构导致的信息泄漏，这使得残差级预训练任务变得微不足道，导致残差表示的表达不足。为了解决这个问题，作者在3D蛋白质链上引入了一个跨度掩模预训练策略，以学习残基和原子的有意义表示。这导致了一种简单而有效的方法来学习适合于各种下游任务的蛋白质表示。结合位点预测和功能预测任务的大量实验结果表明，该预训练方法明显优于其他方法。

蛋白质建模是必不可少的，因为蛋白质在各种细胞过程中起着关键作用，如转录、翻译、信号传导和细胞周期调节。近年来，深度学习的进步极大地促进了用于生成高质量蛋白质表示的预训练模型的发展，从而能够预测蛋白质分类和功能等多种特性。尽管之前的许多研究都集中在基于蛋白质序列的预训练上，蛋白质结构作为蛋白质功能决定因素的重要性导致了基于3D结构模型的出现。最近基于高度精确的深度学习的蛋白质结构预测方法的突破极大地促进了这些模型的发展。

在之前的研究中，一种常见的预训练方法利用了自然语言处理领域中自我预测技术的进步。在这种情况下，给定蛋白质的目标可以被定义为基于剩余结构的信息预测蛋白质的特定片段。标准的预训练任务通常涉及随机屏蔽某些残基，预测被屏蔽残基的类型和位置(即α碳原子的坐标)，以及它们与其他残基之间的角度。通过这一过程，该模型有效地捕获了残基信息，从而获得了蛋白质残基的高质量表征。

尽管许多预先训练的模型已经成功地模拟了蛋白质的3D结构，但大多数模型都集中在残留物水平上，仅利用α碳原子或主链原子的几何位置。然而，侧链原子在许多下游任务中也是必不可少的，例如分子对接，因为它们与小分子相互作用。因此，在蛋白质建模中整合来自所有原子的信息是必要的。在作者的实证研究中，发现寻常的原子级建模通常会失败：(1)简单地将残差输入替换为原子输入，在原子级进行预训练任务，而不考虑残差级别，例如预测原子坐标和角度，并没有产生显著的改进。这表明基于残留物的建模是必不可少的。(2)单纯地将残基和原子信息结合起来，同时在原子和残基层面进行预训练任务，并不能提高性能。

图1：加入全部原子后，残基位置的可能范围受到限制，从而更容易预测残基位置和边角等

作者确定了一个关键原因是由于输入中包含原子结构而导致残差级任务的信息泄漏，这使得残差级预训练任务变得微不足道，并导致残差表示表达不足。如图1所示，当被原子结构包围时，可以很容易地预测当前的残基结构。这说明了对残馀信息进行适当建模的必要性。为了解决这个问题，作者在3D蛋白链上引入了一种跨度掩模策略(SMPC)。作者屏蔽了连续的具有生物学意义的亚结构的残基类型，仅保留了跨屏蔽残基的α碳原子，同时消除了所有其他原子。这种方法增加了残基任务的难度，使得仅从侧链和主链原子推断残基类型和结构变得不可能，从而促使模型学习有意义的残基表示。

模型架构

图2：Vabs-Net架构的概述

如图2所示，作者提出了一个向量感知双层稀疏注意网络(Vabs-Net)，这是一个同时对残基和原子建模的预训练模型。Vabs-Net采用精心设计的边缘矢量编码模块和由原子-原子轨道和残基-残基轨道组成的双轨道稀疏关注模块对原子和残基进行编码。这些轨迹通过碳原子相互作用。为了确保残差水平上的任务是有意义的，作者采用了SMPC预训练策略。在原子水平上，作者采用随机噪声策略。通过一系列的结构预训练任务，如位置和扭转角预测，作者的模型有效地学习残基和原子表示，从而在两个层面上实现全面的蛋白质建模。

残基水平图由所有α碳节点和残基边组成，原子水平图由所有原子节点和原子水平边组成，其中残基水平图与原子水平图共享α碳节点。该网络引入一个虚拟原点，连接每个原子和残基以整合蛋白质的整体表示。对于节点编码而言，节点嵌入结合了原子类型和残基类型，使用大规模语言模型（如ESM）来利用序列信息。对于键距离编码而言，使用高斯核编码原子或残基之间的距离。

传统的蛋白质预训练模型使用残基之间的距离来编码结构信息。然而，对于原子建模来说单纯的距离编码信息不足。本文通过在残基局部坐标系和绝对全局坐标系中编码边的方向来解决这个问题。通过构建一个局部坐标系的旋转矩阵R，用于将全局坐标系中的向量转换为局部坐标系中的向量。

实验结果

表1：EC编号和GO术语预测任务下的模型性能

作者模型在BP、MF、CC和EC任务上的结果见表1，表明Vabs-Net模型在这些任务上的性能显著优于所有基线模型。与GearNet-ESM相比的显著提升证明了原子级别编码的有效性。超过Siamdiff-ESM的优异表现显示了残基级别编码的重要性以及所提出的SMPC预训练策略在学习高效残基级别表示方面的有效性。为了进一步验证SMPC的有效性，作者比较了在预训练过程中有和没有整合SMPC的模型。结果表明，整合SMPC有助于有效的残基级别表示学习。

表2：小分子结合位点的IOU预测结果

表3：DNA结合位点预测结果

表4：RNA结合位点预测结果

表2、表3和表4提供了预训练模型和未预训练模型在结合位点预测上的结果。作者发现，Vabs-Net在所有基线模型中表现最佳。结合位点预测的基线模型通常没有经过预训练，因此作者将未预训练的Vabs-Net与表现最好的基线模型进行了比较。结果显示即使在未预训练的情况下，Vabs-Net的性能也能与最好的基线模型竞争，这证明了该骨干网络的有效性。此外，结合预训练后，性能显著提升，远超其他预训练模型。对比其他预训练模型和未预训练的最佳基线模型，显示了方向编码的重要性。

表5：分子与等效键的对接结果

表5展示了从各种结构蛋白预训练模型中提取的特征在分子对接任务中的效果。结果表明，作者的模型获得的原子和残基表示优于其他预训练模型。通过结合原子级别编码，模型表现出比GearNet更优异的性能。该模型和Siamdiff都优于GearNet进一步强调了原子级别编码的重要性。由于SMPC预训练方法有助于学习更精细的残基级别表示，Vabs-Net模型在性能上超过了Siamdiff。

表6：在小分子结合位点预测任务和EC预测任务中的消融研究

为了分析不同组件的效果，作者选择了蛋白质功能预测任务（EC）和结合位点预测任务进行消融研究。作者调查了不同的预训练策略和模型配置，其结果如表6所示。

1. 对比模型No.0和No.1，展示了原子级别编码的重要性，特别是在侧链原子起关键作用的结合位点预测任务中。
2. 对比模型No.0和No.2，清楚地表明残基级别编码同样重要，因为它增强了感受野并捕捉到残基级别的表示。
3. 对比模型No.0和No.3，确认了SMPC预训练方法在增强残基级别表示方面的有效性。
4. 在去除向量边编码器后（No.3和No.4），性能显著下降，强调了方向编码的重要性，而不仅仅是距离编码。
5. 尽管利用ESM可以提升性能（No.5和No.3），但其影响不如向量边编码器显著。
6. 增加KNN参数并未显著提升模型性能（No.5和No.3），这可能是因为KNN=30足以捕捉大多数结构模体。

结论

在本研究中，作者介绍了Vabs-Net模型，并采用跨跨度掩码策略对3D蛋白质链进行预训练，旨在学习原子级别表示并改进残基级别表示。作者进行了广泛的实验来评估Vabs-Net模型和跨跨度掩码蛋白质链预训练方法的有效性。结果表明，Vabs-Net模型表现优异，超过了之前的最先进模型。在后续研究中，作者计划加强序列和结构特征的整合。

参考资料

Zhao J, Zhuang W, Song J, et al. Multi-level protein pre-training with Vabs-Net[J]. arXiv preprint arXiv:2402.01481, 2024.