【原】【深度学习】Eigenfold：基于diffusion model的生成式蛋白质结构预测模型

——简介——

基于序列预测蛋白质结构的深度学习模型已经取得了极大的成功（如AlphaFold2等结构预测模型），但之前的此类模型只能生成单个蛋白结构。之后有基于生成模型对蛋白构象系综进行预测的工作，但都还处于比较粗糙的阶段。如foldingDiff[1]，仅通过简单地对主链二面角进行扩散进而采样结构。2023年4月，来自MIT CSAIL研究所Tommi Jaakkola课题组的Bowen Jing等人开发了基于diffusion model的从序列生成蛋白质结构的Eigenfold模型，该工作最终发表在ICLR2023 的MLDD workshop上[2]。该模型在精度上接近了AlphaFold2等非生成式模型，并提出了一种新的扩散过程Harmonic Diffusion。

——方法——

Eigenfold取名“Eigen”，在于其提出的一种基于弹性势能的扩散过程Harmonic Diffusion，如下：

其中：

是以单个残基为单位的弹性势能，其中，α=3/3.8Å-2。

这种扩散过程相较于常规的扩散过程最大的好处在于可以避免一些极端不符合物理规律的构象的产生。例如，在之前的很多基于diffusion model的蛋白质结构生成模型中，由于在加噪时直接生搬硬套CV领域的做法，导致在去躁路径中的中间构象往往都是处于极度不符合物理规律的构象。而在eigenfold的去躁路径中产生的构象，往往还保持一些物理上的合理性，如图1所示：

图1 Eigenfold的结构生成过程

对于一般的扩散模型，正向扩散过程可以表示为：

其中，H是表示x各个维度扩散方向和强度的对称半正定矩阵。H可以作如下分解：

其中，P是正交矩阵，Λ是对角线为λ₁至λ_3n的非负对角矩阵。其中，H矩阵的特征向量也即P矩阵的列向量代表扩散的模，Λ的对角线值代表对应模下的强度。同时，diffusion kernal p_t|0和稳态分布p_∞可以沿着这些模变成不相关的高斯函数。之后，p_t|0和稳态分布p_∞的KL散度可以如下表示，同样可以表示为各个模方向上的KL散度加和：

可见，对于不同的扩散强度λ_i，扩散达到收敛的步数会差很远。从另一个角度来说，可以在某些λ_i对应的维度进行扩散时认为具有远大于此时λ_i的λ值的其他维度近似不变（即类似于物理学中的波恩-奥本海默近似）。

基于此，为了更加高效的进行正向扩散和反向推断，作者引入了图像扩散领域使用的级联扩散方法。具体的实现方法为：选定某些截断值τ，当有λ_it>τ时，则将这些维度的扩散坐标设为0，只扩散λ_it<τ的维度。具体到结构上的效果如图2所示，可见在推断过程中，首先是大体的折叠模式确定，之后才是局部主链的环境微调：

图2 对某蛋白结构的扩散推断

至于预测模型，作者将蛋白结构表示成以氨基酸残基为单位的图，包含点特征和边特征。通过训练e3nn网络来进行去躁预测。去躁后的特征通过嵌入到omegafold的structure module中以实现蛋白结构的折叠。

此外，由于模型本身是生成模型，作者为了和AlphaFold等非生成模型进行比较，利用计算最大似然下界的方法对生成的结构进行排序，如下：

具有最大值的结构被认为是最优结构，并进入到后续的benchmark中。

——结果——

作者使用CAMEO数据集（选取750残基以下的数据），将Eigenfold和当下主流的非生成式结构预测模型进行了比较，发现Eigenfold虽然精度最差，但已经比较接近RoseTTAFold的水平，如表1所示：

表1 各个模型的表现（左值为平均值，右值为中位数）

此外，作者比较了omegafold输出的IDDT以及计算的最大似然下界ELBO的回归关系，说明ELBO在一定程度上能够评价生成结构的好坏，如图3所示：

图3  ELBO与IDDT的回归关系

最后，作者评价了模型对于具有多个构象的序列的预测能力，主要通过3个方面来评价：1.在结构全局层面，模型是否能同时采样到两种结构2.样本多样性水平能否预测构象变化的大小？3.残基水平的采样方差和残基实际柔性是否相关？

作者通过计算TM_(con1/con2)来表示两个真实构象之间的结构差异；通过计算TMens来评价第1方面的问题，如下：

其中，x₁，x₂分别表示两个真实构象，y_i表示模型采样构象。

此外，定义TM_var为采样的构象两两之间的TM_score的平均值。最后通过评价发现，模型具有一定的捕捉多种构象的能力，但不多。如图4所示：

图4  模型对结构多样性的采样能力

——小结——

总的来说，Eigenfold是一次利用diffusion model对基于蛋白序列预测蛋白构象系综任务的一次尝试，相比之前工作，笔者认为主要的价值有3点：

1.将生成式模型应用到结构预测任务上，标志着当前AI4S这一领域的任务从单结构预测逐步转向构象系综预测。

2.相较于之前的foldingDiff，将生成模型的预测能力提高了很多。

3.提出了一种基于弹性势能的扩散过程，相较于之前直接高斯加躁的结构生成模型，此模型的去躁路径更具有物理意义。

同时，作者自己也提到，Eigenfold本身也是一项不完美和不成熟的工作，它的后续改进可能会更具有价值，笔者在此想到几点改进的方向：

1.对Embedding的模型进行更换或微调。在本工作中，作者直接使用了omegafold进行特征嵌入，并且没有对模型参数进行微调。

2.对训练数据集进行扩增。单纯使用PDB数据集想要对结构变化进行采样本身就具有困难。实际上，现在已有一些工作（如DiG[3]等），引入MD了的数据对训练数据集进行补充。

参考文献：

[1]   Wu, Kevin E. et al. “Protein structure generation via folding diffusion.” ArXiv abs/2209.15611 (2022): n. pag.

[2]   Jing, Bowen et al. “EigenFold: Generative Protein Structure Prediction with Diffusion Models.” ArXiv (2023): n. pag.

[3]   https://www.microsoft.com/en-us/research/blog/distributional-graphormer-toward-equilibrium-distribution-prediction-for-molecular-systems/