|
——簡(jiǎn)介—— 如何高效采樣蛋白質(zhì)的多種結(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í)的構(gòu)象生成算法被提出，但大都直接生成離散的構(gòu)象系綜，無(wú)法生成具有時(shí)間順序的構(gòu)象變化軌跡。2024年9月，MIT CSAIL研究所Tommi Jaakkola課題組的Bowen Jing等人開(kāi)發(fā)了基于流匹配（flow matching）的基于關(guān)鍵幀的分子動(dòng)力學(xué)模擬軌跡生成模型MDGen [1]。該模型可以基于給定的單一構(gòu)象生成類似于分子動(dòng)力學(xué)模擬軌跡的具有時(shí)間信息的連續(xù)變化構(gòu)象系綜，同時(shí)也能對(duì)于動(dòng)態(tài)軌跡進(jìn)行結(jié)構(gòu)補(bǔ)全（即inpainting）。 ——方法—— MDGen利用3000條三肽和1900多條單鏈蛋白的分子動(dòng)力學(xué)模擬軌跡（各100 ns）進(jìn)行訓(xùn)練，其關(guān)鍵在于對(duì)分子動(dòng)力學(xué)模擬軌跡的特征工程，以便利用神經(jīng)網(wǎng)絡(luò)模型進(jìn)行建模，基本邏輯如下圖1所示。類似于已有的大量蛋白質(zhì)結(jié)構(gòu)深度學(xué)習(xí)模型，MDGen將單個(gè)殘基作為token，但在sequence維度之外增加一個(gè)時(shí)間維度。作者并沒(méi)有直接使用蛋白質(zhì)SE(3)模型常用的residue frame進(jìn)行殘基建模，而是使用殘基相對(duì)于關(guān)鍵幀的位移和旋轉(zhuǎn)信息進(jìn)行建模，具體的tokenize方法如下所示： 
在給定K個(gè)關(guān)鍵幀的情況下，每個(gè)殘基使用7K+14維的特征向量表示。同時(shí)，作者使用了目前生成模型領(lǐng)域較為常用的流匹配模型，具體的網(wǎng)絡(luò)架構(gòu)使用了在時(shí)序深度學(xué)習(xí)領(lǐng)域表現(xiàn)較好的Scalable Interpolant Transformer（SiT）網(wǎng)絡(luò)架構(gòu)。考慮到MD軌跡中的幀數(shù)可能很大（如100 ns的軌跡按10 ps一幀進(jìn)行保存，則會(huì)產(chǎn)生10,000幀），作者將SiT架構(gòu)中的時(shí)序attention部分更換為Hyena架構(gòu)，用于以更低計(jì)算復(fù)雜度代替?zhèn)鹘y(tǒng)的Transformer。作者將訓(xùn)練好的模型應(yīng)用在了4個(gè)下游任務(wù)上：前向模擬（Forward Simulation）、路徑預(yù)測(cè)（Interpolation）、軌跡插幀（Upsampling）、動(dòng)態(tài)修補(bǔ)（Inpainting）。其具體原理如圖1所示：
圖1. MDGen的應(yīng)用場(chǎng)景（左）和實(shí)現(xiàn)方法（右）作者隨后對(duì)MDGen的四種下游任務(wù)進(jìn)行了測(cè)試。對(duì)于Forward Simulation任務(wù)，作者使用第一幀作為參考，訓(xùn)練模型向后生成10 ns的軌跡。在生成時(shí)，模型對(duì)每個(gè)肽段迭代生成100 ns的軌跡，并和真實(shí)軌跡進(jìn)行比較。作者比較的標(biāo)準(zhǔn)主要有：（1）生成軌跡與MD軌跡的結(jié)構(gòu)分布是否匹配，（2）是否能捕捉動(dòng)態(tài)內(nèi)容，（3）是否能在更短的時(shí)間內(nèi)遍歷采樣空間。作者采用對(duì)扭轉(zhuǎn)角和TICA的降維分量計(jì)算Jensen-Shannon散度（JSD）的方法來(lái)衡量結(jié)構(gòu)分布的匹配程度，結(jié)果如圖2所示：
圖 2. Forward Simulation的采樣結(jié)果。（A）MD（橙）與模型生成（藍(lán)）的6種扭轉(zhuǎn)角的分布。（B、C）TICA降維得到的自由能景觀分布。（D）通過(guò)MSM構(gòu)建的亞穩(wěn)態(tài)相關(guān)性。（E）沿TICA的最慢獨(dú)立分量的實(shí)際去相關(guān)時(shí)間。（F）扭轉(zhuǎn)角的去相關(guān)時(shí)間相關(guān)性（藍(lán)），不同亞穩(wěn)態(tài)對(duì)之間的通量矩陣相關(guān)性（橙）。（G）MD和模型生成的丙氨酸四肽構(gòu)象對(duì)比。作者通過(guò)馬爾可夫狀態(tài)模型（MSM）分別得到了真實(shí)軌跡和生成軌跡的亞穩(wěn)態(tài)分布，并發(fā)現(xiàn)模型很好地捕捉到了狀態(tài)的相對(duì)排序，且很少遺漏重要狀態(tài)或?qū)ο∮袪顟B(tài)賦予較高的權(quán)重。此外，通過(guò)比較軌跡中每個(gè)扭轉(zhuǎn)角的去相關(guān)時(shí)間，可以發(fā)現(xiàn)生成軌跡和MD軌跡具有較高的一致性（圖2F）。作者認(rèn)為這一結(jié)果說(shuō)明了模型能夠區(qū)分緩慢和快速弛豫的扭轉(zhuǎn)勢(shì)壘。最后，作者表明模型能夠在約60秒的GPU時(shí)間內(nèi)生成等效于100 ns的軌跡，而MD模擬則需要約3小時(shí)的GPU時(shí)間，強(qiáng)調(diào)了模型生成相比于MD模擬的速度優(yōu)勢(shì)。 對(duì)于給定前后構(gòu)象的變構(gòu)過(guò)程采樣任務(wù)（interpolation / transition path sampling），作者在訓(xùn)練時(shí)給定前后構(gòu)象為關(guān)鍵幀，并使模型產(chǎn)生中間的1 ns構(gòu)象。在測(cè)試時(shí)，作者挑選軌跡中兩個(gè)肽段分離最遠(yuǎn)的狀態(tài)，并產(chǎn)生1000個(gè)中間構(gòu)象。最終的評(píng)估結(jié)果如圖3所示：
圖 3. interpolation的測(cè)試結(jié)果作者對(duì)生成的軌跡和MD軌跡分別構(gòu)建MSM亞穩(wěn)態(tài)，并將100 ns的軌跡MSM亞穩(wěn)態(tài)結(jié)果作為參考標(biāo)準(zhǔn)。如圖3所示，生成的路徑比從任何少于100 ns的重復(fù)MD MSM中采樣的路徑具有更高的可能性，而100 ns正是參考MD模擬的長(zhǎng)度。此外，MDGen的路徑集合在訪問(wèn)狀態(tài)的分布上與參考MD MSM的JSD表現(xiàn)最佳，并且具有最高比例的有效非零概率路徑。 對(duì)于軌跡插幀（upsampling）任務(wù)，作者將10 ps間隔的軌跡插幀至100 fs。通過(guò)比較扭轉(zhuǎn)角的自相關(guān)函數(shù)：
的分布（該函數(shù)自變量為間隔時(shí)間Δt），可以評(píng)估的模型性能，具體如圖4所示：
圖 4. 左圖：扭轉(zhuǎn)角的自相關(guān)函數(shù)（直線代表100 fs-timestep MD，小點(diǎn)代表100 fs-timestep生成結(jié)果，大點(diǎn)代表下采樣的10 ps-timestep結(jié)果）。右圖：自相關(guān)函數(shù)對(duì)對(duì)數(shù)時(shí)間尺度的負(fù)導(dǎo)數(shù)。作者隨后探討了基于動(dòng)力學(xué)軌跡的補(bǔ)全任務(wù)，即在軌跡中遮去一部分結(jié)構(gòu)，讓模型進(jìn)行修補(bǔ)（inpainting）。不同于常規(guī)的蛋白設(shè)計(jì)，作者評(píng)估了針對(duì)四肽（遮去中間兩個(gè)殘基，保留末端殘基）的toy model，在給定觀測(cè)到的殘基構(gòu)象轉(zhuǎn)變的情況下，設(shè)計(jì)在相應(yīng)馬爾可夫狀態(tài)的肽段殘基類型，結(jié)果如表1所示：
表 1. MDGen和baseline模型在序列恢復(fù)率上的比較此外，作者展示了MDGen在較大蛋白的系綜生成任務(wù)上的效果，并和已有的系綜生成模型AlphaFlow進(jìn)行比較。在此任務(wù)中，作者使用ATLAS數(shù)據(jù)集的軌跡進(jìn)行針對(duì)單鏈蛋白的訓(xùn)練。此時(shí)由于殘基數(shù)目大大增加，作者生成250幀、時(shí)間步為400 ps的樣本，使單個(gè)樣本能夠模擬100 ns的ATLAS參考軌跡，具體生成性能如表2：
表 2. MDGen在單鏈蛋白上的結(jié)構(gòu)系綜生成性能 可見(jiàn)，MDGen精確度介于AlphaFlow和MSA下采樣（MSA subsampling）之間，但每生成一個(gè)結(jié)構(gòu)的速度比二者快幾個(gè)數(shù)量級(jí)。圖5是一個(gè)案例蛋白的MD軌跡系綜與MDGen生成系綜的可視化結(jié)果，其中，相似的RMSF圖像也說(shuō)明MDGen除了能較好的復(fù)現(xiàn)結(jié)構(gòu)，也能較好的復(fù)現(xiàn)結(jié)構(gòu)的某些動(dòng)力學(xué)性質(zhì)。
圖 5. （上）MDGen生成的6UOF_A蛋白和MD模擬的系綜（下）RMSF數(shù)據(jù) 相比已有的蛋白構(gòu)象系綜生成模型，MDGen能夠直接生成MD軌跡。但就目前而言，MDGen只能對(duì)短軌跡的小肽進(jìn)行生成，同時(shí)僅對(duì)扭轉(zhuǎn)角、降維數(shù)據(jù)等評(píng)價(jià)指標(biāo)進(jìn)行了分析，尚未提到生成的構(gòu)象在能量上是否合理，是否存在一些原子碰撞等。同時(shí)，作者在文中提到，MDGen受限于關(guān)鍵幀，無(wú)法直接從序列開(kāi)始采樣，且依賴關(guān)鍵幀的準(zhǔn)確度。筆者認(rèn)為MDGen的主要貢獻(xiàn)如下：1. 將類似于視頻生成的訓(xùn)練流程和網(wǎng)絡(luò)架構(gòu)思路引入到MD軌跡中，并很好的和結(jié)構(gòu)深度學(xué)習(xí)結(jié)合起來(lái)。2. 提出了4種下游任務(wù)，啟發(fā)了這類模型的應(yīng)用前景。3. 此類任務(wù)此前缺乏很好的評(píng)測(cè)流程，作者提出的很多評(píng)測(cè)項(xiàng)目都具有相當(dāng)程度的創(chuàng)新性與合理性。[1] Jing B, St?rk H, Jaakkola T, et al. Generative Modeling of Molecular Dynamics Trajectories[J]. arXiv preprint arXiv:2409.17808, 2024.
|
