ICML 2024 | 通过力引导的SE(3)扩散模型生成蛋白质构象

DRUGAI

今天为大家介绍的是来自字节跳动Quanquan Gu团队的一篇论文。蛋白质的构象景观对于理解其在复杂生物过程中的功能至关重要。传统的基于物理的计算方法，如分子动力学（MD）模拟，存在罕见事件采样和长时间平衡问题，限制了它们在一般蛋白质系统中的应用。最近，深度生成建模技术，特别是扩散模型，已被用于生成新颖的蛋白质构象。然而，现有的基于评分的扩散方法无法正确结合重要的物理先验知识来指导生成过程，导致采样的蛋白质构象与平衡分布存在较大偏差。为了解决这些问题，本文提出了一种用于蛋白质构象生成的力引导SE(3)扩散模型——CONFDIFF。通过将力引导网络与基于数据的评分模型混合，CONFDIFF可以生成具有丰富多样性且保持高保真的蛋白质构象。在包括12种快速折叠蛋白质和牛胰岛素抑制剂（BPTI）在内的多种蛋白质构象预测任务上的实验表明，作者的方法优于当前最先进的方法。

蛋白质是动态的大分子，在各种生物过程中发挥着关键作用。它们的功能主要通过构象变化实现，这些结构变化使蛋白质能够与其他分子相互作用。描绘蛋白质的构象景观提供了重要的见解，包括：（1）识别隐藏在蛋白质表面下的潜在药物靶点，以及（2）揭示多种亚稳态之间的过渡路径。全面了解蛋白质构象有助于阐明生物反应机制，从而使研究人员能够设计出具有更高特异性和有效性的靶向抑制剂和治疗剂。

传统的基于物理的模拟方法，如分子动力学（MD）模拟，已被广泛研究用于蛋白质构象采样。然而，为了保持能量守恒和确保数值稳定性，MD模拟的时间步长通常只有几飞秒。因为某些感兴趣的生物过程，如蛋白质折叠，跨越的时间尺度更长，从微秒到秒不等。这导致传统MD模拟的采样效率有限，加上罕见事件采样问题，进一步阻碍了研究界广泛采用MD进行高通量研究。在强大的折叠模型（如AlphaFold、RoseTTAFold、OmegaFold等）的基础上，已经有多次尝试将这些深度神经网络定制用于蛋白质构象采样。然而，这种启发式方法无法保证预测的结构是目标序列的低能状态。最近，一些工作将扩散模型应用于蛋白质构象生成，但现有的扩散模型在利用重要的物理先验信息（如MD力场）来指导其扩散过程方面存在不足，影响了其真实采样符合玻尔兹曼分布的多样蛋白质构象的能力。

为了解决上述挑战，作者提出了一种新颖的力引导扩散模型CONFDIFF，旨在生成更好地符合玻尔兹曼分布的高保真蛋白质构象。受对比能量预测（CEP）技术的启发，作者采用MD能量先验作为基于物理的偏好函数。通过在扩散采样过程中引入额外的力引导网络，优先生成具有较低势能的构象，从而有效提高采样质量。

模型部分

图 1

作者的基线模型包括一个无条件评分模型和一个序列条件评分模型。无条件模型在没有任何序列信息的情况下训练于蛋白质结构，有效地捕捉一般蛋白质的构象分布。一方面，序列条件模型能够访问蛋白质序列信息（seq）以及时间t对应的结构。

作者采用与FramePred相似的网络架构来参数化相应的评分函数。无条件模型将残基索引和扩散时间t的正弦嵌入作为其单一（{}）和对（{}）嵌入，而条件模型则额外将ESMFold预计算的表征拼接到其单一嵌入中。需要注意的是，条件模型的序列表示选择是灵活的，因为已有研究表明，使用预训练折叠模型的表示有助于扩散模型生成合理的蛋白质结构，而无条件模型则可以有效提高采样多样性。两个模型都使用去噪评分匹配（DSM）损失函数进行训练，其中，λ(t) 是一个与评分范数成反比的重新加权函数。

在反向采样过程中，作者使用超参数γ来控制条件模型的无分类器引导强度，以便通过下式估计评分函数。

在蛋白质构象建模的背景下，可以访问基于物理的能量函数（即原子之间的势能）及其梯度（即每个原子的力）。与非规范化的势能函数相比，原子力更加局部，并且表现出更好的数值稳定性，这也更符合评分匹配的目标。作者使用上面介绍的基线评分模型从生成蛋白质构象，然后使用这些构象训练一个独立的中间力网络。力引导的训练过程总结在算法1中，推理过程总结在算法2中。

算法 1

算法 2

力引导的构象采样

作者将CONFDIFF与两种最先进的基于扩散的蛋白质构象生成模型进行比较：EIGENFOLD、STR2STR。所有基线模型都在其默认设置下运行。对于CONFDIFF，作者按照STR2STR的方法，以不同的序列条件水平（γ）采样构象。所有模型都在蛋白质数据库（PDB）中存储的实验蛋白质结构上进行训练，而不使用额外的MD轨迹数据。

图 2

作为初步概念验证，作者首先展示了提出的力引导如何通过生成能量更低（即稳定性更高）的蛋白质构象来提高CONFDIFF的采样能力。作者选择了Lindorff-Larsen等研究中的快速折叠蛋白质之一WW域作为示例，使用CONFDIFF在不同程度的力引导（η）和序列条件（γ）下生成构象。如图2所示，在没有力引导（η = 0）的情况下，具有较弱序列条件的模型可以生成多样性更高（RMSF）的样本，但同时也具有更高的能量，显示了质量与多样性的权衡。在整合力引导后，CONFDIFF减轻了这种影响，能够生成能量更低且RMSF水平相似的样本。这个结果表明，力引导可以在不显著降低多样性的情况下提高构象的稳定性。这种有利的效果表明，整合物理信息可能为局部优化结构的采样提供细致的指导。

快速折叠蛋白质的分布预测

该数据集包含12种短蛋白质，在此实验中，作者评估模型恢复模拟中观察到的构象分布的能力。为了评估分布相似性，作者计算生成样本分布与真实MD样本分布之间的JS距离，并评估有效性得分和RMSF以反映多样性。此外，作者报告预测的残基间接触率的均方根误差（RMSEcontact），以反映模型在柔性区域的准确性。

表 1

图 3

结果总结在表1中，并在图3和图S3中展示TIC投影中的样本分布。CONFDIFF在恢复样本分布方面一贯优于EIGENFOLD和STR2STR，表现为更低的JS距离和更低的残基间接触预测均方根误差（RMSE）。值得注意的是，整合能量和力引导提高了样本的有效性，同时保持了类似的样本多样性，证实了整合物理引导的益处。相比之下，EIGENFOLD显示出有限的样本多样性，这可能是由于作为完全条件扩散模型强烈倾向于预测折叠状态。STR2STR利用无条件扩散模型来探索多样的结构空间并生成多样的样本。然而，CONFDIFF更好的分布相关评分表明，通过序列条件控制的扩散过程可能生成更符合蛋白质真实分布的样本。

BPTI的稳态预测

先前对BPTI的MD模拟研究恢复了5个动能簇，代表了其天然折叠状态附近的亚稳态。对于这个基准，作者评估模型恢复这5种状态的质量和效率。与快速折叠蛋白质相比，这个基准要求模型生成不同的构象，同时保持正确的折叠结构。

表 2

作者通过样本到簇中心的最低RMSD来衡量模型的精度。为了比较，作者报告5个簇的平均最佳RMSD（RMSDAVG）和簇3的RMSD（RMSDCLS3），后者是最难采样的簇。如表2所示，带有力引导的CONFDIFF在这两个指标上表现最佳，表明其在预测不同亚稳态方面具有更好的能力。

编译 | 黄海涛

审稿 | 曾全晨

参考资料

Wang, Y., Wang, L., Shen, Y., Wang, Y., Yuan, H., Wu, Y., & Gu, Q. (2024). Protein Conformation Generation via Force-Guided SE (3) Diffusion Models. arXiv preprint arXiv:2403.14088.