文 | 千芯科技董事長(zhǎng) 陳巍
就在前幾天，迄今為止參數(shù)最多、規(guī)模最大的蛋白質(zhì)預(yù)測(cè)模型ESMFold被Meta官宣了，甚至有研究者宣稱(chēng)該模型又大又好，足以碾壓Google在2021年推出的AlphaFold2。
▲ESMFold與通訊作者M(jìn)eta AI的Alexander
這一消息著實(shí)讓學(xué)術(shù)界和工業(yè)界震撼，要知道這些大的模型，無(wú)論訓(xùn)練還是使用，都得有妥妥的“鈔能力”，如果模型越來(lái)越小，說(shuō)不定就不需要更大算力的芯片了。（當(dāng)然事實(shí)并非如此）甚至LeCun大牛都發(fā)推為ESMFold背書(shū)，稱(chēng)之為“Super-fast and accurate”。
從氨基酸序列預(yù)測(cè)蛋白質(zhì)結(jié)構(gòu)是自然科學(xué)中長(zhǎng)期存在的重大挑戰(zhàn)。在基于進(jìn)化的算法中，AlphaFold2可以說(shuō)是目前解決該問(wèn)題最成功的。它通過(guò)在多序列輸入、進(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ì)專(zhuān)注于大規(guī)模蛋白質(zhì)序列和結(jié)構(gòu)數(shù)據(jù)的無(wú)監(jiān)督表示學(xué)習(xí)模型研究。Alexander本人同時(shí)也是Fate Therapeutics、Syros Pharma、Kallyope的聯(lián)合創(chuàng)始人，妥妥的科創(chuàng)家。
那ESMFold真的能碾壓AlphaFold2嗎？讓我們先來(lái)回顧下什么是蛋白質(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)。由線(xiàn)性氨基酸組成的蛋白質(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ì)多肽鏈的線(xiàn)性氨基酸序列。
蛋白質(zhì)二級(jí)結(jié)構(gòu)：依靠不同氨基酸之間的C=O和N-H基團(tuán)間的氫鍵形成的穩(wěn)定結(jié)構(gòu)，主要為α螺旋和β折疊。
蛋白質(zhì)三級(jí)結(jié)構(gòu)：通過(guò)多個(gè)二級(jí)結(jié)構(gòu)元素在三維空間的排列所形成的一個(gè)蛋白質(zhì)分子的三維結(jié)構(gòu)。
蛋白質(zhì)四級(jí)結(jié)構(gòu)：用于描述由不同多肽鏈（亞基）間相互作用形成具有功能的蛋白質(zhì)復(fù)合物分子。
我們所說(shuō)的蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)（Protein Structure Prediction），就是指從蛋白質(zhì)的氨基酸序列中預(yù)測(cè)蛋白質(zhì)的三維結(jié)構(gòu)。也就是說(shuō)，從蛋白質(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萬(wàn)的蛋白質(zhì)結(jié)構(gòu)。
▲冷凍電鏡及其圖像
因此，使用AI的方法，加速對(duì)蛋白質(zhì)結(jié)構(gòu)的解析，分析其組成和功能，就成了生物界和醫(yī)藥界的爭(zhēng)相推進(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)空間。
過(guò)去，AlphaFold2和RoseTTAFold在原子分辨率蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)問(wèn)題上取得了突破性成功，但依賴(lài)于使用多序列比對(duì)（Multiple Sequence Alignment，簡(jiǎn)寫(xiě)為MSA）和相似蛋白質(zhì)結(jié)構(gòu)的模板來(lái)實(shí)現(xiàn)最優(yōu)表現(xiàn)。
▲ESMFold模型具有比AlphaFold2更高的速度
ESMFold使用ESM-2學(xué)習(xí)的信息和表示來(lái)執(zhí)行端到端的3D結(jié)構(gòu)預(yù)測(cè)，特別是僅使用單個(gè)序列作為輸入（AlphaFold2需要多序列輸入），方便研究者在使用時(shí)通過(guò)模型縮放，將模型大小控制在數(shù)百萬(wàn)到數(shù)十億量級(jí)參數(shù)。需要注意的是，隨著模型大小的增加，可觀(guān)察到預(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ù)庫(kù)。使用ESMFold，可以快速計(jì)算100萬(wàn)個(gè)預(yù)測(cè)結(jié)構(gòu)，這些結(jié)構(gòu)代表了蛋白質(zhì)預(yù)測(cè)空間的不同子集，其中大多數(shù)沒(méi)有注釋的結(jié)構(gòu)或功能。
而且ESMFold的大部分高置信度預(yù)測(cè)與已知的實(shí)驗(yàn)結(jié)構(gòu)的相似度都很低，這表明了通過(guò)AI計(jì)算獲得的基因組蛋白的結(jié)構(gòu)新穎性。
值得注意的是，許多高置信度結(jié)構(gòu)與UniRef90中的結(jié)構(gòu)也具有低序列相似性，說(shuō)明該模型具有超出其訓(xùn)練數(shù)據(jù)集的泛化能力，實(shí)現(xiàn)了基于結(jié)構(gòu)的蛋白質(zhì)功能預(yù)見(jiàn)能力。
據(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模型類(lèi)似，ESMFold模型的架構(gòu)也可以分為四部分：數(shù)據(jù)解析部分、編碼器部分（Folding Trunk）、解碼器部分（Structure Module）、循環(huán)部分（Recycling）。
ESMFold和AlphaFold2之間的一個(gè)關(guān)鍵區(qū)別是使用語(yǔ)言模型表示來(lái)消除對(duì)顯式同源序列（以MSA的形式）作為輸入的要求。
語(yǔ)言模型表示作為輸入提供給ESMFold的折疊主干。通過(guò)將處理MSA的計(jì)算量大的Folding Block模塊替換為處理序列的Tranformer模塊來(lái)簡(jiǎn)化AlphaFold2中的Evoformer。這種簡(jiǎn)化或優(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)行，無(wú)需訪(fǎng)問(wèn)任何外部數(shù)據(jù)庫(kù)。
例如在單個(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ù)庫(kù)的初始版本發(fā)布時(shí)具有約36萬(wàn)個(gè)預(yù)測(cè)結(jié)構(gòu)，截至2022年7月則包含約99.5萬(wàn)個(gè)預(yù)測(cè)，這比目前許多蛋白質(zhì)序列數(shù)據(jù)庫(kù)小幾個(gè)數(shù)量級(jí)。
04.數(shù)據(jù)解析部分與解碼器的深度分析
數(shù)據(jù)解析部分用于輸入序列和數(shù)據(jù)庫(kù)的解析，為編碼器提供輸入。
在AlphaFold2模型中，數(shù)據(jù)解析部分使用了氨基酸序列數(shù)據(jù)庫(kù)和結(jié)構(gòu)數(shù)據(jù)庫(kù)，分別用于相近序列的比對(duì)和結(jié)構(gòu)模板的配對(duì)。
▲AlphaFold2多序列比對(duì)示意
生物信息學(xué)的基礎(chǔ)是基于這樣的一個(gè)假設(shè)：序列相似，結(jié)構(gòu)相似，功能相似。一般認(rèn)為相近的序列或者相近的結(jié)構(gòu)會(huì)衍生出相近的功能域。
1）序列數(shù)據(jù)庫(kù)被用于多序列比對(duì)（Multiple Sequence Alignment，MSA），即在序列數(shù)據(jù)庫(kù)中檢索與輸入序列接近的數(shù)據(jù)庫(kù)序列。
2）結(jié)構(gòu)數(shù)據(jù)庫(kù)則用于結(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使用語(yǔ)言模型表示，消除了對(duì)明確的同源序列（以MSA的形式）作為輸入的需要。
ESMFold通過(guò)用一個(gè)處理序列的Transformer模塊取代處理MSA的計(jì)算昂貴的網(wǎng)絡(luò)模塊，簡(jiǎn)化了AlphaFold2中的Evoformer。這種簡(jiǎn)化意味著ESMFold的速度大大提高，遠(yuǎn)高于基于MSA的模型。
05.結(jié)語(yǔ)
作為蛋白質(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削減了搜索多序列來(lái)構(gòu)建MSA的計(jì)算量。
▲ESMFold用于探索宏基因組結(jié)構(gòu)空間
推斷速度優(yōu)勢(shì)使得基于計(jì)算有效映射大型宏基因組序列數(shù)據(jù)庫(kù)的結(jié)構(gòu)空間成為可能。
除了用于識(shí)別遠(yuǎn)同源性外，ESMFold還可以被用于進(jìn)行快速準(zhǔn)確的結(jié)構(gòu)預(yù)測(cè)，并在實(shí)際時(shí)間尺度內(nèi)獲得數(shù)百萬(wàn)個(gè)預(yù)測(cè)結(jié)構(gòu)，進(jìn)一步幫助發(fā)現(xiàn)新的蛋白質(zhì)結(jié)構(gòu)和功能。這相當(dāng)于在使用AI計(jì)算來(lái)構(gòu)建生命的“元宇宙”。
150億參數(shù)大模型，10x倍速度提升。雖然Meta ESMFold精度上沒(méi)能做到全面“碾壓”AlphaFold2，但畢竟“唯快不破”，對(duì)于蛋白質(zhì)結(jié)構(gòu)解析與預(yù)測(cè)、構(gòu)建大型宏基因組結(jié)構(gòu)數(shù)據(jù)庫(kù)有著巨大的推動(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.