文 | 千芯科技董事長 陳巍
就在前幾天,迄今為止參數(shù)最多、規(guī)模最大的蛋白質(zhì)預(yù)測(cè)模型ESMFold被Meta官宣了,甚至有研究者宣稱該模型又大又好,足以碾壓Google在2021年推出的AlphaFold2。
▲ESMFold與通訊作者M(jìn)eta AI的Alexander
這一消息著實(shí)讓學(xué)術(shù)界和工業(yè)界震撼,要知道這些大的模型,無論訓(xùn)練還是使用,都得有妥妥的“鈔能力”,如果模型越來越小,說不定就不需要更大算力的芯片了。(當(dāng)然事實(shí)并非如此)甚至LeCun大牛都發(fā)推為ESMFold背書,稱之為“Super-fast and accurate”。
從氨基酸序列預(yù)測(cè)蛋白質(zhì)結(jié)構(gòu)是自然科學(xué)中長期存在的重大挑戰(zhàn)。在基于進(jìn)化的算法中,AlphaFold2可以說是目前解決該問題最成功的。它通過在多序列輸入、進(jìn)化同源物對(duì)齊序列和可選結(jié)構(gòu)模板上訓(xùn)練端到端神經(jīng)網(wǎng)絡(luò),取得了突破性成就,大大加速了“生命元宇宙”的構(gòu)建。
▲蛋白質(zhì)預(yù)測(cè)AI大模型的進(jìn)化
而Meta團(tuán)隊(duì)的ESMFold蛋白質(zhì)模型只需要一個(gè)序列作為輸入,該模型背后的團(tuán)隊(duì)由Meta AI(原Facebook AI)的資深研究科學(xué)家Alexander Rives主導(dǎo)。該團(tuán)隊(duì)專注于大規(guī)模蛋白質(zhì)序列和結(jié)構(gòu)數(shù)據(jù)的無監(jiān)督表示學(xué)習(xí)模型研究。Alexander本人同時(shí)也是Fate Therapeutics、Syros Pharma、Kallyope的聯(lián)合創(chuàng)始人,妥妥的科創(chuàng)家。
那ESMFold真的能碾壓AlphaFold2嗎?讓我們先來回顧下什么是蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè),然后再深入分析ESMFold的網(wǎng)絡(luò)結(jié)構(gòu)。
▲ESMFold預(yù)測(cè)的結(jié)構(gòu)
論文鏈接:https://doi.org/10.1101/2022.07.20.500902
01.什么是蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)?
首先,蛋白質(zhì)結(jié)構(gòu)是指各種蛋白質(zhì)分子的空間結(jié)構(gòu)。由線性氨基酸組成的蛋白質(zhì),需要折疊(Fold)成特定的空間結(jié)構(gòu),才具有相應(yīng)的生理活性和生物學(xué)功能。
▲蛋白質(zhì)的四級(jí)結(jié)構(gòu)
蛋白質(zhì)的分子結(jié)構(gòu)可劃分為四級(jí),以描述其不同層級(jí)的特征:
蛋白質(zhì)一級(jí)結(jié)構(gòu):組成蛋白質(zhì)多肽鏈的線性氨基酸序列。
蛋白質(zhì)二級(jí)結(jié)構(gòu):依靠不同氨基酸之間的C=O和N-H基團(tuán)間的氫鍵形成的穩(wěn)定結(jié)構(gòu),主要為α螺旋和β折疊。
蛋白質(zhì)三級(jí)結(jié)構(gòu):通過多個(gè)二級(jí)結(jié)構(gòu)元素在三維空間的排列所形成的一個(gè)蛋白質(zhì)分子的三維結(jié)構(gòu)。
蛋白質(zhì)四級(jí)結(jié)構(gòu):用于描述由不同多肽鏈(亞基)間相互作用形成具有功能的蛋白質(zhì)復(fù)合物分子。
我們所說的蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)(Protein Structure Prediction),就是指從蛋白質(zhì)的氨基酸序列中預(yù)測(cè)蛋白質(zhì)的三維結(jié)構(gòu)。也就是說,從蛋白質(zhì)的一級(jí)結(jié)構(gòu)預(yù)測(cè)其折疊和二級(jí)、三級(jí)、四級(jí)結(jié)構(gòu)。
DeepMind(Google旗下)的AlphaFold2在蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)大賽CASP 14中,對(duì)大部分蛋白質(zhì)結(jié)構(gòu)的預(yù)測(cè)與真實(shí)結(jié)構(gòu)只差一個(gè)原子的寬度,達(dá)到接近冷凍電鏡等復(fù)雜儀器檢測(cè)的水平。這一巨大進(jìn)步被Nature和Science選為2021年度十大科學(xué)突破。
根據(jù)不同的氨基酸和序列,蛋白質(zhì)能折疊成的構(gòu)型數(shù)量是一個(gè)天文數(shù)字,因此很難用常規(guī)方法進(jìn)行蛋白質(zhì)結(jié)構(gòu)的準(zhǔn)確預(yù)測(cè)。例如,目前實(shí)驗(yàn)的方法(例如冷凍電鏡)至今才能解出10萬的蛋白質(zhì)結(jié)構(gòu)。
▲冷凍電鏡及其圖像
因此,使用AI的方法,加速對(duì)蛋白質(zhì)結(jié)構(gòu)的解析,分析其組成和功能,就成了生物界和醫(yī)藥界的爭相推進(jìn)的重要工作。
02.ESMFold的“魔幻效果”
ESMFold與AlphaFold2和RoseTTAFold對(duì)多序列輸入的蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)具有相當(dāng)?shù)臏?zhǔn)確度。但ESMFold突出優(yōu)勢(shì)在于,其計(jì)算速度比AlphaFold2快一個(gè)數(shù)量級(jí),能夠在更有效的時(shí)間尺度上探索蛋白質(zhì)的結(jié)構(gòu)空間。
過去,AlphaFold2和RoseTTAFold在原子分辨率蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)問題上取得了突破性成功,但依賴于使用多序列比對(duì)(Multiple Sequence Alignment,簡寫為MSA)和相似蛋白質(zhì)結(jié)構(gòu)的模板來實(shí)現(xiàn)最優(yōu)表現(xiàn)。
▲ESMFold模型具有比AlphaFold2更高的速度
ESMFold使用ESM-2學(xué)習(xí)的信息和表示來執(zhí)行端到端的3D結(jié)構(gòu)預(yù)測(cè),特別是僅使用單個(gè)序列作為輸入(AlphaFold2需要多序列輸入),方便研究者在使用時(shí)通過模型縮放,將模型大小控制在數(shù)百萬到數(shù)十億量級(jí)參數(shù)。需要注意的是,隨著模型大小的增加,可觀察到預(yù)測(cè)準(zhǔn)確性的持續(xù)提升。(還是“越大越準(zhǔn)”)
▲ESM-2模型隨著參數(shù)量升高精度升高
由于ESMFold的預(yù)測(cè)速度比現(xiàn)有的其他原子分辨率蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)模型快一個(gè)數(shù)量級(jí),因此ESMFold可以幫助快速構(gòu)建蛋白質(zhì)結(jié)構(gòu)數(shù)據(jù)庫。使用ESMFold,可以快速計(jì)算100萬個(gè)預(yù)測(cè)結(jié)構(gòu),這些結(jié)構(gòu)代表了蛋白質(zhì)預(yù)測(cè)空間的不同子集,其中大多數(shù)沒有注釋的結(jié)構(gòu)或功能。
而且ESMFold的大部分高置信度預(yù)測(cè)與已知的實(shí)驗(yàn)結(jié)構(gòu)的相似度都很低,這表明了通過AI計(jì)算獲得的基因組蛋白的結(jié)構(gòu)新穎性。
值得注意的是,許多高置信度結(jié)構(gòu)與UniRef90中的結(jié)構(gòu)也具有低序列相似性,說明該模型具有超出其訓(xùn)練數(shù)據(jù)集的泛化能力,實(shí)現(xiàn)了基于結(jié)構(gòu)的蛋白質(zhì)功能預(yù)見能力。
據(jù)此,研究人員認(rèn)為,ESMFold可以幫助理解那些超出現(xiàn)有認(rèn)知的蛋白質(zhì)結(jié)構(gòu)。
▲ESMFold在單序列輸入時(shí)預(yù)測(cè)精度明顯好于AlphaFold2
雖然ESMFold速度很高,精度也不錯(cuò),特別是在單序列輸入的時(shí)候精度明顯好于AlphaFold2。但我們也要看到,ESMFold在多序列輸入的情況下,其精度比AlphaFold2還是略有差距。
03.ESMFold網(wǎng)絡(luò)結(jié)構(gòu)
與AlphaFold2模型類似,ESMFold模型的架構(gòu)也可以分為四部分:數(shù)據(jù)解析部分、編碼器部分(Folding Trunk)、解碼器部分(Structure Module)、循環(huán)部分(Recycling)。
ESMFold和AlphaFold2之間的一個(gè)關(guān)鍵區(qū)別是使用語言模型表示來消除對(duì)顯式同源序列(以MSA的形式)作為輸入的要求。
語言模型表示作為輸入提供給ESMFold的折疊主干。通過將處理MSA的計(jì)算量大的Folding Block模塊替換為處理序列的Tranformer模塊來簡化AlphaFold2中的Evoformer。這種簡化或優(yōu)化意味著ESMFold會(huì)比基于MSA的模型快得多。
▲ESMFold與AlphaFold2對(duì)比
在AlphaFold2和RoseTTAFold中使用MSA和模板會(huì)導(dǎo)致兩個(gè)瓶頸。
首先,可能需要基于CPU檢索和對(duì)齊MSA和模板。這是由于AlphaFold2和RoseTTAFold不是二維序列嵌入狀態(tài),而是使用軸向注意力對(duì)應(yīng)于MSA的三維內(nèi)部狀態(tài)進(jìn)行操作,即使使用GPU,這一計(jì)算的代價(jià)也不菲。
相比之下,ESMFold是一個(gè)完全端到端的序列結(jié)構(gòu)預(yù)測(cè)器,可以完全在GPU上運(yùn)行,無需訪問任何外部數(shù)據(jù)庫。
例如在單個(gè)NVIDIA V100 GPU上,使用較少參數(shù)的ESMFold在14.2秒內(nèi)對(duì)具有384個(gè)殘基的蛋白質(zhì)進(jìn)行預(yù)測(cè),可比單個(gè)AlphaFold2模型快6倍。而在較短的序列上,我們甚至看到了約60倍的改進(jìn)。
速度的數(shù)量級(jí)提高是ESMFold優(yōu)于AlphaFold2的獨(dú)特優(yōu)勢(shì),使我們能夠在比現(xiàn)有方法更短的時(shí)間尺度內(nèi)構(gòu)建大量預(yù)測(cè)結(jié)構(gòu)??紤]到可用序列數(shù)據(jù)的規(guī)模,這一點(diǎn)尤其重要。
例如,AlphaFold2蛋白質(zhì)結(jié)構(gòu)數(shù)據(jù)庫的初始版本發(fā)布時(shí)具有約36萬個(gè)預(yù)測(cè)結(jié)構(gòu),截至2022年7月則包含約99.5萬個(gè)預(yù)測(cè),這比目前許多蛋白質(zhì)序列數(shù)據(jù)庫小幾個(gè)數(shù)量級(jí)。
04.數(shù)據(jù)解析部分與解碼器的深度分析
數(shù)據(jù)解析部分用于輸入序列和數(shù)據(jù)庫的解析,為編碼器提供輸入。
在AlphaFold2模型中,數(shù)據(jù)解析部分使用了氨基酸序列數(shù)據(jù)庫和結(jié)構(gòu)數(shù)據(jù)庫,分別用于相近序列的比對(duì)和結(jié)構(gòu)模板的配對(duì)。
▲AlphaFold2多序列比對(duì)示意
生物信息學(xué)的基礎(chǔ)是基于這樣的一個(gè)假設(shè):序列相似,結(jié)構(gòu)相似,功能相似。一般認(rèn)為相近的序列或者相近的結(jié)構(gòu)會(huì)衍生出相近的功能域。
1)序列數(shù)據(jù)庫被用于多序列比對(duì)(Multiple Sequence Alignment,MSA),即在序列數(shù)據(jù)庫中檢索與輸入序列接近的數(shù)據(jù)庫序列。
2)結(jié)構(gòu)數(shù)據(jù)庫則用于結(jié)構(gòu)匹配,尋找與輸入序列的結(jié)構(gòu)接近的已知結(jié)構(gòu)模板。
然后序列比對(duì)與結(jié)構(gòu)比對(duì)的結(jié)果作為輸入傳輸給編碼器部分。
▲ESMFold Folding Block與AlphaFold2 Evoformer結(jié)構(gòu)對(duì)比
解碼器部分即Folding Trunk,一共48層。
ESMFold與AlphaFold2的一個(gè)關(guān)鍵區(qū)別是,ESMFold使用語言模型表示,消除了對(duì)明確的同源序列(以MSA的形式)作為輸入的需要。
ESMFold通過用一個(gè)處理序列的Transformer模塊取代處理MSA的計(jì)算昂貴的網(wǎng)絡(luò)模塊,簡化了AlphaFold2中的Evoformer。這種簡化意味著ESMFold的速度大大提高,遠(yuǎn)高于基于MSA的模型。
05.結(jié)語
作為蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)大模型,ESMFold獲得準(zhǔn)確原子分辨率結(jié)構(gòu)預(yù)測(cè)的推斷(Inferenc)速度比AlphaFold2提高了約一個(gè)數(shù)量級(jí)。特別是在實(shí)際計(jì)算中,這一速度優(yōu)勢(shì)表現(xiàn)的更加明顯。這是由于ESMFold削減了搜索多序列來構(gòu)建MSA的計(jì)算量。
▲ESMFold用于探索宏基因組結(jié)構(gòu)空間
推斷速度優(yōu)勢(shì)使得基于計(jì)算有效映射大型宏基因組序列數(shù)據(jù)庫的結(jié)構(gòu)空間成為可能。
除了用于識(shí)別遠(yuǎn)同源性外,ESMFold還可以被用于進(jìn)行快速準(zhǔn)確的結(jié)構(gòu)預(yù)測(cè),并在實(shí)際時(shí)間尺度內(nèi)獲得數(shù)百萬個(gè)預(yù)測(cè)結(jié)構(gòu),進(jìn)一步幫助發(fā)現(xiàn)新的蛋白質(zhì)結(jié)構(gòu)和功能。這相當(dāng)于在使用AI計(jì)算來構(gòu)建生命的“元宇宙”。
150億參數(shù)大模型,10x倍速度提升。雖然Meta ESMFold精度上沒能做到全面“碾壓”AlphaFold2,但畢竟“唯快不破”,對(duì)于蛋白質(zhì)結(jié)構(gòu)解析與預(yù)測(cè)、構(gòu)建大型宏基因組結(jié)構(gòu)數(shù)據(jù)庫有著巨大的推動(dòng)作用。
參考文獻(xiàn):
Zeming Lin et. al., Language models of protein sequences at the scale of evolution enable accurate structure prediction, https://www.biorxiv.org/content/10.1101/2022.07.20.500902v1
Jumper, J. et al., Highly accurate protein structure prediction with AlphaFold, Nature (2021):1-11.