Str2Str：基于分?jǐn)?shù)模型的zero-shot蛋白質(zhì)構(gòu)象采樣方法

——簡(jiǎn)介——如何高效采樣蛋白質(zhì)尤其是天然無(wú)規(guī)蛋白質(zhì)（IDP）的多種結(jié)構(gòu)一直是計(jì)算結(jié)構(gòu)生物學(xué)領(lǐng)域的一大難點(diǎn)。通常，此類問(wèn)題是通過(guò)分子模擬（包括MC采樣和MD采樣）來(lái)解決，同時(shí)配合各種增強(qiáng)采樣算法提高效率。近年來(lái)，得益于GPU算力的進(jìn)步，各種基于深度學(xué)習(xí)的增強(qiáng)采樣和構(gòu)象生成算法被提出，但大都需要針對(duì)某一體系預(yù)先的MD軌跡數(shù)據(jù)，沒(méi)有辦法直接從序列生成具有多樣性的構(gòu)象。2024年1月，來(lái)自蒙特利爾大學(xué)Mila研究院JianTang課題組的JiaruiLu等人開發(fā)了基于分?jǐn)?shù)模型蛋白質(zhì)zero-shot構(gòu)象采樣模型Str2str[1]。該工作目前在ICLR2024上作為Poster發(fā)布。該模型通過(guò)輸入蛋白序列和結(jié)構(gòu)（例如使用AlphaFold預(yù)測(cè)的結(jié)構(gòu)），在一個(gè)分?jǐn)?shù)模型上加噪再去噪，從而采樣新的構(gòu)象。該工具不同于之前的AI輔助增強(qiáng)采樣方法和構(gòu)象生成方法，不需要對(duì)結(jié)構(gòu)提前進(jìn)行MD采樣，大大擴(kuò)展了模型的使用范圍?！椒ā猄tr2str全稱structuretostructure，其基本邏輯如下圖1所示：先利用蛋白質(zhì)序列-結(jié)構(gòu)數(shù)據(jù)集訓(xùn)練一個(gè)具有平移旋轉(zhuǎn)等變性的分?jǐn)?shù)模型；之后針對(duì)某一輸入的蛋白序列通過(guò)結(jié)構(gòu)預(yù)測(cè)模型如AlphaFold2等預(yù)測(cè)一個(gè)初始結(jié)構(gòu)；將結(jié)構(gòu)按一定步數(shù)不完全加噪；利用訓(xùn)練好的模型對(duì)結(jié)構(gòu)進(jìn)行去噪；最后得到的結(jié)構(gòu)即為圖1.Str2str的運(yùn)行邏輯這個(gè)先不完全加噪再去噪的過(guò)程被作者稱為forward-backward過(guò)程（FB），具體公式如下：其中，Tδ取值在0到T之間，并且可以通過(guò)控制其大小，來(lái)達(dá)到控制生成構(gòu)象集合（ensemble）的結(jié)構(gòu)在訓(xùn)練時(shí)，分?jǐn)?shù)模型的隨機(jī)微分方程（SDE）如下所示：為了更好的訓(xùn)練效果，作者仿照AlphaFold2的處理方法來(lái)表征蛋白結(jié)構(gòu)：將蛋白質(zhì)按氨基酸為單位劃分為主鏈“frame”以及側(cè)鏈各種扭轉(zhuǎn)角。其中，對(duì)于蛋白主鏈的表征定義如下：其中，Xglobal代表蛋白的笛卡爾坐標(biāo)，Xlocal代表在“frame”參考系下的坐標(biāo)。Tframe可分項(xiàng)Rframe和一個(gè)平移項(xiàng)vframe。而對(duì)于側(cè)鏈，則仿照AF2的側(cè)鏈對(duì)照表處理，本推文暫不贅述。至于去噪的神經(jīng)網(wǎng)絡(luò)架構(gòu)，作者使用了一個(gè)更改后的IPA網(wǎng)絡(luò)，被作者命名為DenoisingIPA，如下圖所圖2.DenoisingIPA的網(wǎng)絡(luò)架構(gòu)可見，基本與AF2中的IPA模塊相同，其中Si張量來(lái)源于positionalencoding和分?jǐn)?shù)模型的Timestepencoding。此外，作者在其中加入了一個(gè)pairrepresentation的transition模塊，類似于AF2的Evoformer中所使用的edgetransition模塊，如下：值得注意的是，不同AF2的IPA模塊，這個(gè)架構(gòu)中其實(shí)并沒(méi)有嵌入序列信息，所以其實(shí)蛋白序列的唯一作用是在一開始通過(guò)AF2或ESMFold等工具預(yù)測(cè)初始結(jié)構(gòu)。在訓(xùn)練過(guò)程中，作者是用了兩種loss函數(shù)來(lái)保證訓(xùn)練效果。其一是正常分?jǐn)?shù)模型的loss函數(shù)：與正常分?jǐn)?shù)模型訓(xùn)練過(guò)程不同的是，由于在實(shí)際使用時(shí)，并不需要用到全部的擴(kuò)散步數(shù)，因此可以只對(duì)0另一種loss被稱為AuxiliaryLoss，由主鏈MSEloss：和原子距離圖損失distogramloss：這兩種loss在t<T/4時(shí)被計(jì)算。最后，總loss如下：其中，作者將alpha和beta取值為0.25?！Y(jié)果——之后，作者對(duì)Str2str的效果進(jìn)行了benchmark。Benchmark的對(duì)象主要包括AF2-RVAE、EigenFold、idpGAN等。評(píng)價(jià)標(biāo)準(zhǔn)主要包括：1.有效性：包括原子clash和化學(xué)鍵鍵長(zhǎng)。2.保真度：兩兩距離分布的JS散度（JS-PwD）和tICA降維圖上的最近組分（JS-TIC）以及回旋半徑Rg分布的JS散度（JS-Rg）。3.多樣性：包括RMSD和TM-score的MAE。Benchmark結(jié)果如下表1所示，其中PF表示使用概率流進(jìn)行采樣的結(jié)果：表1.Str2str的benchmark結(jié)果可見，與之前的EigenFold和idpGAN相比，Str2str不僅在有效性和保真度上有優(yōu)勢(shì)，在多樣性上超過(guò)之前的方法很多，同時(shí)和MD的結(jié)果進(jìn)行比較也能發(fā)現(xiàn)，單從采樣效率上講，Str2str的似乎媲美幾十us級(jí)的除了量化的benchmark之外，作者對(duì)很多MD采樣領(lǐng)域比較著名的case進(jìn)行了具體的分析，包括很多的fast-folding蛋白如：BPTI、Trp-cage圖3.BPTI的benchmark結(jié)果可見，Str2str的生成結(jié)果和長(zhǎng)時(shí)間MD更類似，而其他方法則和短MD的結(jié)果更為類似，體現(xiàn)了Str2str更Trp-cage的contact-map分布圖如下圖所示：圖4.Trp-cage的benchmark結(jié)果其他各種case的contact-map結(jié)果如下如所示：圖5.各種fast-folding蛋白的benchmark結(jié)果此外，作者還統(tǒng)計(jì)了在不同去噪步數(shù)的選取情況下Str2str的表現(xiàn)，展現(xiàn)了模型可以通過(guò)參數(shù)控制生成結(jié)果多樣性的特點(diǎn)，如下圖所示：圖6.不同T_delta下的ensemble生成結(jié)果比較有意思的一點(diǎn)是隨著T_delta的增大,模型似乎能更容易采樣到不同的構(gòu)象。下圖的tICA結(jié)果也說(shuō)明了圖7.不同T_delta下的WWdomainensemble生成結(jié)果和MD結(jié)果比較——小結(jié)——Str2str的提出不僅為我們提供了一種高效的AI采樣方法，同時(shí)這種不完全加噪方式也為將擴(kuò)散模型遷移到flow上提供了一些啟示。此外，據(jù)作者言，這種不完全加噪再去噪的過(guò)程實(shí)際上暗合MC采樣或MD采樣中跨越勢(shì)能壘的過(guò)程，同樣引人深思。圖8.Str2str與常規(guī)采樣方法的能壘跨越過(guò)程類比當(dāng)然，就目前而言基于AI的增強(qiáng)采樣方法仍然任重道遠(yuǎn)：比如，Str2str的