比AlphaFold2快一個數(shù)量級！蛋白質(zhì)通用大模型來了，13個任務(wù)取得SOTA丨百圖生科&清華

蕭簫 2023-07-08 13:09:01 來源：量子位

蕭簫 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI

蛋白質(zhì)領(lǐng)域的“通用大模型”來了！

參數(shù)量高達(dá)1000億，在這個領(lǐng)域尚屬首次，一出場就在13個任務(wù)上達(dá)到SOTA——

例如在抗體結(jié)構(gòu)預(yù)測任務(wù)中，就超越了“老前輩”AlphaFold2。

如果說之前蛋白質(zhì)AI模型還大多停留在單一類型任務(wù)上，現(xiàn)在從蛋白質(zhì)預(yù)測到蛋白質(zhì)設(shè)計，各種細(xì)分任務(wù)都用一個AI就能完成。

這個蛋白質(zhì)語言模型由百圖生科和清華大學(xué)共同開發(fā)，取名xTrimoPGLM。

所以，它究竟在哪些任務(wù)上取得了SOTA，這個千億大模型又究竟是如何煉成的？背后的運(yùn)作原理和實(shí)現(xiàn)方式是什么？

未來在整個生命科學(xué)領(lǐng)域，是否也會像自然語言這樣，出現(xiàn)類似ChatGPT的通用大模型？

百圖生科的CTO宋樂博士向我們分享了思考。

蛋白質(zhì)的“通用大模型”長啥樣？

這個名叫xTrimoPGLM的蛋白質(zhì)語言大模型，“底子”是清華推出的GLM（通用語言模型）算法。

選用這種算法，是因?yàn)槿祟愓Z言和蛋白質(zhì)之間存在很多相似之處。

和語言任務(wù)一樣，蛋白質(zhì)任務(wù)也可以被分為理解（預(yù)測）和生成（設(shè)計）兩大類：

生成任務(wù)，指根據(jù)不同條件設(shè)計對應(yīng)的蛋白質(zhì)。如給定某個功能標(biāo)簽，要求生成能實(shí)現(xiàn)這一功能的蛋白質(zhì)；或是給一段蛋白質(zhì)結(jié)構(gòu)，設(shè)計一段可折疊成該結(jié)構(gòu)的氨基酸序列等。

理解任務(wù)，指預(yù)測某種蛋白質(zhì)的屬性。如蛋白質(zhì)本身是酶，用AI預(yù)測它的最優(yōu)催化溫度、催化效率、穩(wěn)定性等。

但和人類語言一樣，在大模型出現(xiàn)前，蛋白質(zhì)語言模型往往也“只能干好一件事”，同時學(xué)多個任務(wù)，反而可能把原來的能力給“忘了”。

（畢竟理解類任務(wù)通常用的是雙向注意力機(jī)制，類似完形填空；但生成類任務(wù)用的卻是單向的，像續(xù)寫作文一樣）

而GLM在框架設(shè)計上兼顧了大模型的理解和生成能力，因此也成為團(tuán)隊這次的基礎(chǔ)架構(gòu)“首選項(xiàng)”。

不過，相比自然語言模型，蛋白質(zhì)在訓(xùn)練數(shù)據(jù)、任務(wù)和框架適用性上又有所不同。

為了更好地將GLM的優(yōu)勢和蛋白質(zhì)語言特性結(jié)合起來，xTrimoPGLM設(shè)計了一種新框架，其中增加了MLM（掩碼語言模型）部分。

其中，紫色的[MASK]表示MLM，用于提升模型理解能力；綠色的[sMASK]和[gMASK]表示GLM，用于提升模型生成能力。

具體到細(xì)節(jié)上，[sMASK]掩蓋序列中間的連續(xù)部分，模型預(yù)測時不僅要學(xué)會預(yù)測內(nèi)容，還需要學(xué)會預(yù)測長度；[gMASK]掩蓋除了上下文之外的其余序列部分，以進(jìn)一步提升模型的生成能力。

不過，即使是Meta的蛋白質(zhì)理解模型ESM，參數(shù)量也就在150億級別左右。

為何xTrimoPGLM模型參數(shù)量會達(dá)到千億級？

宋樂博士介紹稱，這是因?yàn)榈鞍踪|(zhì)的數(shù)據(jù)比想象中要更大：

目前可用的蛋白質(zhì)序列已經(jīng)有幾十億，而這些序列的長度平均又達(dá)到幾百甚至上千，乘起來就接近自然語言token的數(shù)量規(guī)模了。

更大的數(shù)據(jù)量，自然需要更大的模型來“吃”。

在大語言模型已經(jīng)達(dá)到千億級參數(shù)量的情況下，蛋白質(zhì)語言模型理論上也應(yīng)達(dá)到這一規(guī)模，才能實(shí)現(xiàn)比Meta的ESM更好的效果。

基于這一理念設(shè)計的xTrimoPGLM，確實(shí)在理解和生成任務(wù)上均取得了不錯的效果。

斬下13個SOTA，可直接用于行業(yè)

研究人員一共將xTrimoPGLM在15個任務(wù)上進(jìn)行了測試。

事實(shí)證明，這個蛋白質(zhì)語言模型在其中的13個任務(wù)上都取得了SOTA。

這些任務(wù)從蛋白質(zhì)結(jié)構(gòu)、可發(fā)展性、相互作用到功能分為四大類，具體又包括評估蛋白質(zhì)特性，如溶解性、對蛋白酶的穩(wěn)定性、溫度敏感性、蛋白質(zhì)結(jié)合親和力、抗生素抗性等。

以抗體結(jié)構(gòu)預(yù)測為例。據(jù)宋樂博士介紹，和AlphaFold2相比，xTrimoPGLM不僅效果更好，而且速度快了接近一個數(shù)量級。

之所以能做到在模型更大的同時，預(yù)測速度還更快了，是因?yàn)橄啾華lphaFold2，xTrimoPGLM“跳了一步”：

AlphaFold2依賴多序列比對進(jìn)行搜索，但xTrimoPGLM因?yàn)橐呀?jīng)“學(xué)會了蛋白質(zhì)的語言”，所以直接省去了這個步驟。

在此基礎(chǔ)之上，xTrimoPGLM不僅能很好地提供蛋白質(zhì)序列信息，模型能力也得到了增強(qiáng)。

宋樂博士認(rèn)為，不止是抗體結(jié)構(gòu)預(yù)測，類似思路也能被推廣到更通用的蛋白質(zhì)結(jié)構(gòu)預(yù)測上去，這也在團(tuán)隊的下一步計劃之內(nèi)。

不僅如此，團(tuán)隊還計劃把模型擴(kuò)展到RNA、DNA等不同模態(tài)的生命科學(xué)數(shù)據(jù)上，甚至是跨細(xì)胞、跨組織層面，嘗試實(shí)現(xiàn)更加通用的生命科學(xué)大模型。

當(dāng)然，目前這一階段的xTrimoPGLM，就已經(jīng)能直接提供給行業(yè)使用了。

從它能實(shí)現(xiàn)的任務(wù)來看，已經(jīng)涵蓋了不少蛋白質(zhì)下游應(yīng)用的場景，如涉及蛋白質(zhì)的相關(guān)酶的設(shè)計，以及醫(yī)藥食品行業(yè)的一些消費(fèi)級蛋白質(zhì)預(yù)測任務(wù)，都可以直接用xTrimoPGLM去幫助解決。

據(jù)宋樂博士介紹，xTrimoPGLM未來也會接入到百圖生科的AI生成蛋白平臺AIGP中去，負(fù)責(zé)如抗體結(jié)構(gòu)預(yù)測、親和力預(yù)測和蛋白質(zhì)相互作用預(yù)測等任務(wù)。

One More Thing

目前來看，發(fā)展“通用大模型”的路徑主要有兩種。

一種是繼續(xù)擴(kuò)大單個模型的參數(shù)量，試圖達(dá)到真正的單個AGI之路；

另一個則是通過多個模型聯(lián)合的如Mixture of Experts等方式，將負(fù)責(zé)不同任務(wù)類型的大模型之間聯(lián)合起來，以實(shí)現(xiàn)更多功能。

這兩條路各有其優(yōu)缺點(diǎn)所在。如果要想繼續(xù)發(fā)展生命科學(xué)領(lǐng)域的“通用大模型”，哪條路徑更有可能通往AGI？

宋樂博士認(rèn)為“都有可能”。

不過就百圖生科團(tuán)隊而言，他們?nèi)匀徊扇《鄠€大模型聯(lián)動的方式，來繼續(xù)探索通用大模型之路。

論文地址：
https://www.biorxiv.org/content/10.1101/2023.07.05.547496v1