文 | 千芯科技董事長 陳巍

就在前幾天，迄今為止參數(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.