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