【原】【深度學(xué)習(xí)】Eigenfold：基于diffusion model的生成式蛋白質(zhì)結(jié)構(gòu)預(yù)測模型

GoDesign 2023-08-23 發(fā)布于北京

展開全文

——簡介——

基于序列預(yù)測蛋白質(zhì)結(jié)構(gòu)的深度學(xué)習(xí)模型已經(jīng)取得了極大的成功（如AlphaFold2等結(jié)構(gòu)預(yù)測模型），但之前的此類模型只能生成單個蛋白結(jié)構(gòu)。之后有基于生成模型對蛋白構(gòu)象系綜進(jìn)行預(yù)測的工作，但都還處于比較粗糙的階段。如foldingDiff[1]，僅通過簡單地對主鏈二面角進(jìn)行擴散進(jìn)而采樣結(jié)構(gòu)。2023年4月，來自MIT CSAIL研究所Tommi Jaakkola課題組的Bowen Jing等人開發(fā)了基于diffusion model的從序列生成蛋白質(zhì)結(jié)構(gòu)的Eigenfold模型，該工作最終發(fā)表在ICLR2023 的MLDD workshop上[2]。該模型在精度上接近了AlphaFold2等非生成式模型，并提出了一種新的擴散過程Harmonic Diffusion。

——方法——

Eigenfold取名“Eigen”，在于其提出的一種基于彈性勢能的擴散過程Harmonic Diffusion，如下：

其中：

是以單個殘基為單位的彈性勢能，其中，α=3/3.8?-2。

這種擴散過程相較于常規(guī)的擴散過程最大的好處在于可以避免一些極端不符合物理規(guī)律的構(gòu)象的產(chǎn)生。例如，在之前的很多基于diffusion model的蛋白質(zhì)結(jié)構(gòu)生成模型中，由于在加噪時直接生搬硬套CV領(lǐng)域的做法，導(dǎo)致在去躁路徑中的中間構(gòu)象往往都是處于極度不符合物理規(guī)律的構(gòu)象。而在eigenfold的去躁路徑中產(chǎn)生的構(gòu)象，往往還保持一些物理上的合理性，如圖1所示：

圖1 Eigenfold的結(jié)構(gòu)生成過程

對于一般的擴散模型，正向擴散過程可以表示為：

其中，H是表示x各個維度擴散方向和強度的對稱半正定矩陣。H可以作如下分解：

其中，P是正交矩陣，Λ是對角線為λ₁至λ_3n的非負(fù)對角矩陣。其中，H矩陣的特征向量也即P矩陣的列向量代表擴散的模，Λ的對角線值代表對應(yīng)模下的強度。同時，diffusion kernal p_t|0和穩(wěn)態(tài)分布p_∞可以沿著這些模變成不相關(guān)的高斯函數(shù)。之后，p_t|0和穩(wěn)態(tài)分布p_∞的KL散度可以如下表示，同樣可以表示為各個模方向上的KL散度加和：

可見，對于不同的擴散強度λ_i，擴散達(dá)到收斂的步數(shù)會差很遠(yuǎn)。從另一個角度來說，可以在某些λ_i對應(yīng)的維度進(jìn)行擴散時認(rèn)為具有遠(yuǎn)大于此時λ_i的λ值的其他維度近似不變（即類似于物理學(xué)中的波恩-奧本海默近似）。

基于此，為了更加高效的進(jìn)行正向擴散和反向推斷，作者引入了圖像擴散領(lǐng)域使用的級聯(lián)擴散方法。具體的實現(xiàn)方法為：選定某些截斷值τ，當(dāng)有λ_it>τ時，則將這些維度的擴散坐標(biāo)設(shè)為0，只擴散λ_it<τ的維度。具體到結(jié)構(gòu)上的效果如圖2所示，可見在推斷過程中，首先是大體的折疊模式確定，之后才是局部主鏈的環(huán)境微調(diào)：

圖2 對某蛋白結(jié)構(gòu)的擴散推斷

至于預(yù)測模型，作者將蛋白結(jié)構(gòu)表示成以氨基酸殘基為單位的圖，包含點特征和邊特征。通過訓(xùn)練e3nn網(wǎng)絡(luò)來進(jìn)行去躁預(yù)測。去躁后的特征通過嵌入到omegafold的structure module中以實現(xiàn)蛋白結(jié)構(gòu)的折疊。

此外，由于模型本身是生成模型，作者為了和AlphaFold等非生成模型進(jìn)行比較，利用計算最大似然下界的方法對生成的結(jié)構(gòu)進(jìn)行排序，如下：

具有最大值的結(jié)構(gòu)被認(rèn)為是最優(yōu)結(jié)構(gòu)，并進(jìn)入到后續(xù)的benchmark中。

——結(jié)果——

作者使用CAMEO數(shù)據(jù)集（選取750殘基以下的數(shù)據(jù)），將Eigenfold和當(dāng)下主流的非生成式結(jié)構(gòu)預(yù)測模型進(jìn)行了比較，發(fā)現(xiàn)Eigenfold雖然精度最差，但已經(jīng)比較接近RoseTTAFold的水平，如表1所示：

表1 各個模型的表現(xiàn)（左值為平均值，右值為中位數(shù)）

此外，作者比較了omegafold輸出的IDDT以及計算的最大似然下界ELBO的回歸關(guān)系，說明ELBO在一定程度上能夠評價生成結(jié)構(gòu)的好壞，如圖3所示：

圖3 ELBO與IDDT的回歸關(guān)系

最后，作者評價了模型對于具有多個構(gòu)象的序列的預(yù)測能力，主要通過3個方面來評價：1.在結(jié)構(gòu)全局層面，模型是否能同時采樣到兩種結(jié)構(gòu)2.樣本多樣性水平能否預(yù)測構(gòu)象變化的大小？3.殘基水平的采樣方差和殘基實際柔性是否相關(guān)？

作者通過計算TM_(con1/con2)來表示兩個真實構(gòu)象之間的結(jié)構(gòu)差異；通過計算TMens來評價第1方面的問題，如下：

其中，x₁，x₂分別表示兩個真實構(gòu)象，y_i表示模型采樣構(gòu)象。

此外，定義TM_var為采樣的構(gòu)象兩兩之間的TM_score的平均值。最后通過評價發(fā)現(xiàn)，模型具有一定的捕捉多種構(gòu)象的能力，但不多。如圖4所示：

圖4 模型對結(jié)構(gòu)多樣性的采樣能力

——小結(jié)——

總的來說，Eigenfold是一次利用diffusion model對基于蛋白序列預(yù)測蛋白構(gòu)象系綜任務(wù)的一次嘗試，相比之前工作，筆者認(rèn)為主要的價值有3點：

1.將生成式模型應(yīng)用到結(jié)構(gòu)預(yù)測任務(wù)上，標(biāo)志著當(dāng)前AI4S這一領(lǐng)域的任務(wù)從單結(jié)構(gòu)預(yù)測逐步轉(zhuǎn)向構(gòu)象系綜預(yù)測。

2.相較于之前的foldingDiff，將生成模型的預(yù)測能力提高了很多。

3.提出了一種基于彈性勢能的擴散過程，相較于之前直接高斯加躁的結(jié)構(gòu)生成模型，此模型的去躁路徑更具有物理意義。

同時，作者自己也提到，Eigenfold本身也是一項不完美和不成熟的工作，它的后續(xù)改進(jìn)可能會更具有價值，筆者在此想到幾點改進(jìn)的方向：

1.對Embedding的模型進(jìn)行更換或微調(diào)。在本工作中，作者直接使用了omegafold進(jìn)行特征嵌入，并且沒有對模型參數(shù)進(jìn)行微調(diào)。

2.對訓(xùn)練數(shù)據(jù)集進(jìn)行擴增。單純使用PDB數(shù)據(jù)集想要對結(jié)構(gòu)變化進(jìn)行采樣本身就具有困難。實際上，現(xiàn)在已有一些工作（如DiG[3]等），引入MD了的數(shù)據(jù)對訓(xùn)練數(shù)據(jù)集進(jìn)行補充。

參考文獻(xiàn)：

[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/