ICLR 2023 | 迈向高效有效的蛋白质反向折叠

编译 | 曾全晨 审稿 | 王建民

今天为大家介绍的是来自Stan Z. Li团队的一篇蛋白质设计论文。如何有效且高效地设计能折叠成所需结构的蛋白质序列呢？近年来，基于结构的蛋白质设计中的人工智能方法引起了越来越多的关注；然而，由于缺乏富有表现力的特征和自回归序列解码器，很少有方法能够同时提高准确性和效率。为了解决这些问题，作者提出了PiFold方法，它包含了一种新颖的残基特征提取器和PiGNN层，以改进回复过程，能够以一次性的方式生成蛋白质序列。实验证明，PiFold在CATH 4.2数据集上的恢复率达到了51.66%，而推理速度比自回归方法快70倍。

图 1

蛋白质是由氨基酸组成的线性链，通过折叠成3D结构来控制细胞过程，例如转录、翻译、信号传导和细胞周期调控。为了人类目的创造新的蛋白质可以加深我们对生命系统的理解，并促进对抗疾病的斗争。其中一个关键问题是设计能够折叠成所需结构的蛋白质序列，即基于结构的蛋白质设计。近年来，许多深度学习模型已经被提出来解决这个问题，其中图形模型取得了重要进展。然而，模型的准确性和效率仍有提升的空间。例如，由于缺乏富有表现力的残基表示，大多数图形模型在CATH数据集上无法达到50%以上的序列恢复率。此外，它们还受到自回归生成的限制，导致推理速度较慢。作者的目标是用一个简单的模型，尽可能少地包含冗余，同时提高准确性和效率。多年来，基于图形的模型通过更好的特征工程、更精细的模型和更大的训练数据集努力学习富有表现力的残基表示。尽管目前的模型取得了显著的进展，但在相同的数据集设置下仍有两个问题需要解决：（1）有没有更好的方法来构建有效的特征以促进学习残基表示？（2）我们如何改进模型，使其能够从残基相互作用中学习更好的表示？

大多数图形模型采用自回归解码方案来生成氨基酸，从而显著减慢了推理过程。很少有研究尝试提高模型的效率，也许是因为效率的提升需要牺牲一些准确性，而在蛋白质设计中，准确性比效率更重要。为了解决这个困境，AlphaDesign提出了一个并行的自校正模块，以加速推理过程，同时几乎保持恢复率不变。然而，这仍然会导致一些性能下降，并且需要进行两次迭代进行预测。作者提出了PiFold来解决上述问题，其中包含了一种新颖的残基特征提取器和堆叠的PiGNN。对于特征提取器，针对每个残基，作者构建了更全面的特征，并引入了可学习的虚拟原子来捕捉真实原子所忽略的信息。PiGNN在节点、边和全局级别考虑特征之间的依赖关系，以从多尺度的残基相互作用中学习。此外，可以通过堆叠更多的PiGNN层完全去除自回归解码器，而不损失准确性。实验证明，PiFold在几个真实世界数据集上实现了最先进的恢复率。

模型架构

图 2

图2中展示了PiFold的整体框架，其中输入是蛋白质结构，输出是预期能够折叠成输入结构的蛋白质序列。提出了一种新颖的残基特征提取器和PiGNN层，用于学习具有表达能力的残基表示。具体而言，残基特征提取器构建了全面的残基特征，并创建了可学习的虚拟原子来捕捉与实际原子不同的信息。PiGNN在节点、边缘和全局上下文层面考虑了多尺度的残基相互作用。PiFold能够以一次生成的方式生成蛋白质序列，并且具有比先前的自回归方法更高的恢复率。

图 3

图3描述了如何构建蛋白质图和残基特征。图形结构：作者将蛋白质表示为从残基派生的k-NN图，以考虑三维依赖关系，其中k默认为30。蛋白质图由邻接矩阵，节点特征和边特征组成。节点特征和边特征考虑残基的稳定结构、顺序和坐标的归纳偏置。由于每个残基由C、N和O组成，可以为残基构造局部坐标系。基于这些坐标系，可以为单个残基或残基对构造旋转和平移不变的特征，包括距离、角度和方向。

作者提出了PiGNN层，以考虑节点、边和全局上下文层次中的多尺度残基相互作用，从而学习几何残基表示。对于局部节点交互，作者使用简化的graph transformer来更新节点表示。蛋白质图是一个带有节点和边特征的属性图；然而，简化的graph transformer是一个以节点为中心的网络，它不会迭代地更新边特征。作者发现忽略边特征的更新会导致次优表示，因此引入了边更新层可以改善模型的能力。尽管局部相互作用在学习残基表示中起着至关重要的作用，但全局信息也被证明对改进蛋白质设计具有价值。然而，全局注意力跨整个蛋白质的时间复杂度与蛋白质长度的平方成正比，这会显著增加计算开销。为了同时获得改进的恢复率和良好的效率，作者建议为每个蛋白质学习一个全局上下文向量，并将其用于门控注意力模块来更新节点特征。

实验部分

表 1

作者将PiFold与最近的强基线模型在CATH数据集上进行了比较。作者使用与GraphTrans和GVP相同的数据分割方法，将蛋白质按照CATH拓扑分类进行分区，其中训练集包含18024个蛋白质，验证集包含608个蛋白质，测试集包含1120个蛋白质。表1中报告了困惑度和恢复得分。"短链"数据集包含长度不超过100的蛋白质，"单链"数据集包含在蛋白质数据库中记录为单链的蛋白质。可以观察到，所提出的PiFold能够在不同的测试集上持续改善困惑度（较低为更好）和恢复得分（较高为更好）。在"短链"和"全部"数据集上，PiFold实现了最佳的困惑度和恢复得分，其中恢复率分别提高了3.49%和1.94%。在"单链"数据集上，当在CATH 4.2.4上训练时，PiFold是最佳模型。

图 4

我们已经知道PiFold可以实现更高的恢复率；这种改进是以效率为代价吗？进一步评估了PiFold和竞争基线模型（AlphaDesign和ProteinMPNN）的训练和推断时间成本。在训练阶段，PiFold能够以更少的训练轮数实现最先进的困惑度和恢复率。如图4所示，PiFold只需20轮就足够。至于推断速度，除了AlphaDesign之外，其他基线模型均采用自回归解码器逐个生成残基，计算复杂度与蛋白质长度L成正比，即O(L)。然而，由于PiFold采用一次性生成模式，其计算复杂度为O(1)。在具有平均长度为1632的100个长链上对不同的图模型进行了推断速度的基准测试，并在图1中展示了结果，其中PiFold的速度比自回归竞争对手（包括ProteinMPNN、ESM-IF、GCA、StructGNN、StructTrans、GVP）快70倍。

结论

作为一种基于结构的蛋白质设计的人工智能方法，PiFold通过使用提出的蛋白质特征化器和PiGNN，显著提高了恢复率并大大提高了效率。它在测试集上分别达到了非常高的恢复得分，并且在设计长蛋白质的推断速度上比自回归竞争对手快70倍。

参考资料

Gao, Z., Tan, C., & Li, S. Z. (2022). PiFold: Toward effective and efficient protein inverse folding. arXiv preprint arXiv:2209.12643.