arXiv｜字节跳动通过力引导的SE(3)扩散模型生成蛋白质构象

蛋白质的构象景观对于理解它们在复杂生物过程中的功能至关重要。传统的基于物理的计算方法，如分子动力学(MD)模拟，存在罕见事件采样和长平衡时间的问题，阻碍了其在一般蛋白质建模中的应用。近年来，深度生成模型技术，特别是扩散模型，已被用于生成新的蛋白质构象。然而，现有的基于分数的扩散模型方法不能适当地纳入重要的物理先验知识来引导生成过程，导致采样的蛋白质构象与平衡分布之间存在较大偏差。

为了克服这些限制，本文提出了一种力引导的SE(3) (Euclidean 3D transformation) 扩散模型CONFDIFF，用于蛋白质构象生成。CONFDIFF通过将力引导网络与基于数据的评分模型相结合，可以生成丰富多样的蛋白质构象，同时保持高保真度。在包括12种蛋白质折叠和牛胰蛋白酶抑制剂(BPTI)的实验中，CONFDIFF超越了最先进的方法。

相关研究以Protein Conformation Generation via Force-Guided SE(3) Diffusion Models为题，2024年3月21日发表在arXiv上。

研究背景

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

传统的基于物理的模拟方法，如分子动力学模拟，在蛋白质构象采样中得到了广泛的研究。该类模型通过结合精心设计的经验力场和数值打分器，基于牛顿力学原理，能够模拟蛋白质系统三维结构随时间的变化。在足够的时间内，MD模拟收敛于平衡分布（即玻尔兹曼分布），这有助于估计重要的热力学性质，例如结合自由能变化。然而，为了节省能量和保证数值稳定性，MD模拟的时间步长通常只有几个飞秒，这带来了挑战。因为重要的生物过程，如蛋白质折叠，跨越更长的时间尺度，从微秒到秒。这导致传统MD模拟的采样效率有限，再加上罕见事件采样问题，阻碍了研究界广泛采用MD进行高通量研究。现已有多种神经网络用于蛋白质构象采样，诸如AlphaFold，RoseTTAFold，OmegaFold等。然而，这种启发式方法不能保证预测的结构是目标序列的低能态。近期也有一些研究纳入了扩散模型生成蛋白质构象。然而，现有的扩散模型在利用重要的物理先验信息（如MD力场）来指导其扩散过程方面存在不足，阻碍了其采样符合玻尔兹曼分布的多种蛋白质构象的能力。本文提出的新型力引导扩散模型CONFDIFF，旨在生成更符合玻尔兹曼分布的高保真蛋白构象。从对比能量预测(CEP)技术中获得灵感，CONFDIFF采用MD能量先验作为基于物理的偏好函数。通过在扩散采样过程中引入额外的力引导网络，优先生成势能较低的构象，有效地提高了采样质量。

研究内容

CONFDIFF的主要贡献突出如下:

1. 采用序列条件模型来指导无条件模型，在SE(3)上使用无分类器指导，以在构象质量和多样性之间找到更好的平衡。与DiG方法相比，CONFDIFF在训练过程中不依赖MD数据;与STR2STR方法相比，引导强度系数为平衡样本多样性和质量提供了更高的自由度。

2. CONFDIFF利用MD能量函数作为基于物理的奖励来指导蛋白质构象的生成。此外，在扩散采样过程中采用中间力引导策略。这是第一个已知适用于蛋白质构象生成的力引导网络，有助于将各种构象预测与平衡分布对齐。

3. 在各种基准测试上的实验表明，CONFDIFF方法优于最先进的方法。特别是，能量和力的引导有效地将模型引导到能量较低构象的样本上，使得不同的样本更符合底层的玻尔兹曼分布。

图1 CONFDIFF示意图

方法

本文评估了不同指导策略下CONFDIFF在蛋白质构象生成任务中的有效性，将CONFDIFF与两种最先进的基于扩散的蛋白质构象生成模型EIGENFOLD、STR2STR进行了比较，并将CONFDIFF应用于快速蛋白质折叠和牛胰腺胰蛋白酶抑制剂(BPTI)这两种蛋白质构象基准进行全面评估。

1. 力引导的构象采样

在快速折叠蛋白中（WW Domain蛋白为例），使用CONFDIFF在不同程度的力引导(η)和序列条件(γ)下生成构象结果如图2所示。结果表明，力引导可以在不显著降低多样性的情况下提高构象稳定性。

图2 在不同的力引导(η)和序列条件(γ)下，WW Domain的采样构象的能量(左)和多样性(右)

2. 快速折叠蛋白的分布预测

研究者使用CONFDIFF评估了模型恢复模拟中观察到的构象分布的能力。评估了生成样本和真实MD样本分布之间的JS距离(Jensen-Shannon distance)，以及diversity的validity得分和RMSF(Root Mean Square Fluctuation)，并报告了残基之间预测接触率RMSEcontact(Root Mean Square Error)，以反映柔性区域的模型精度。表1为结果，并在图3中展示了TIC(time-lagged independent components)投影中的样本分布（折叠态的实验结构以颜色表示,来自参考MD的5个随机样本显示为灰色）。在预测残基残基接触时，CONFDIFF始终优于EIGENFOLD和STR2STR，均具有较低的JS距离和较低的RMSE。这说明整合能量和力引导提高了CONFDIFF-base的样本有效性，并保持了相似的样本多样性。较高的CONFDIFF分布相关分数表明，由序列条件控制的扩散过程可能产生更接近蛋白质真实分布的构象。

表1 快速折叠蛋白的结果

图3 WW Domain前两个TIC components的样本分布

3. BPTI的亚稳态预测

本文评估了模型恢复5种BPTI原生折叠态附近亚稳态的质量和效率。与快速折叠蛋白质相比，该基准要求模型在保持正确折叠结构的同时产生不同的构象。指标为5个聚类的最佳RMSD平均值(RMSDAVG) 和聚类3(最难采样的集群)的RMSD平均值(RMSDCLS3)。如表2所示，带力引导的CONFDIFF在这两个指标上都表现最好，这表明CONFDIFF在预测不同亚稳态方面有更好的能力。

表2 BPTI的5个亚稳态恢复情况

本文进一步通过比较不同样本量下的精度来评估模型的效率, 如图4A所示（参考结构用颜色表示，样本结构用灰色表示）。所有CONFDIFF模型对聚类3的采样效率都很好，而EIGENFOLD由于样本多样性低，很快就趋于平稳。CONFDIFF-force表现出最好的性能。

图4 BPTI的亚稳态预测。A)预测聚类3与不同样本量的精度。B)聚类1和聚类3的三种模型最佳采样视觉比较

总结

本文通过混合序列条件评分网络和无条件评分模型提出CONFDIFF来生成蛋白质构象。CONFDIFF通过无分类器的引导在采样质量和多样性之间取得了平衡。在此基础上，作者提出了新的基于物理的能量和力引导策略，并提供了理论保证，有效地引导扩散采样器产生更符合底层玻尔兹曼分布的低能量构象。在各种蛋白质系统上的实验证明了所提出的方法的有效性。

虽然CONFDIFF已经显示出良好的性能，但仍有进一步改进的空间。由于现有蛋白质结构数据库的构象多样性有限，完全以序列信息为条件的模型往往只能预测折叠结构，缺乏全面描绘整个构象景观的能力。另一方面，在生成式建模框架下使用全原子MD能量函数用作偏好函数仍然是相对计算密集和耗时的。然而，这是第一次尝试将玻尔兹曼先验与扩散模型相结合，本文作者将继续追求更有效的力引导方法，以及在训练过程中最小化方差的有效措施。

参考文献

[1] Wang Y, Wang L, Shen Y, et al. Protein Conformation Generation via Force-Guided SE (3) Diffusion Models. arxiv preprint arxiv:2403.14088, 2024.

--------- End ---------