128張GPU煉出中國(guó)版AlphaFold2：訓(xùn)練代碼全開源，打開瀏覽器就可體驗(yàn)，北大元培系A(chǔ)I公司打造

楊凈 夢(mèng)晨 發(fā)自 凹非寺

量子位 報(bào)道 | 公眾號(hào) QbitAI

中國(guó)版AlphaFold2，來(lái)了！

Uni-Fold，官宣即開源，還為科研者提供了測(cè)試接口和在線試用。

半年前，AlphaFold2和RoseTTAFold兩個(gè)頂級(jí)AI算法齊齊引發(fā)生物學(xué)界、AI學(xué)界大地震。

半年后，這個(gè)橫空出世的Uni-Fold，不光能完整復(fù)現(xiàn)AlphaFold2。

其預(yù)測(cè)精度超過(guò)RoseTTAFold、與AlphaFold2相當(dāng)，效率還超AlphaFold2?2-3倍。

△Uni-Fold預(yù)測(cè)的蛋白結(jié)構(gòu)

更重要的是，背后的團(tuán)隊(duì)不是如谷歌這樣的互聯(lián)網(wǎng)巨頭，也不是數(shù)十年深耕于此的頂尖高校團(tuán)隊(duì)。

而是一家成立剛3年的AI創(chuàng)業(yè)公司——深勢(shì)科技。

這個(gè)Uni-Fold究竟有何看頭？接下來(lái)一看究竟。

128張GPU煉出中國(guó)版AlphaFold2

AlphaFold2開源的消息曾一度火爆全場(chǎng)，等稍微冷靜下來(lái)大家才發(fā)現(xiàn)，真正把它用起來(lái)并不容易。

用谷歌官方提供的Colab資源嗎？可以，請(qǐng)排隊(duì)。

自己有算力資源想部署？也可以，不過(guò)AlphaFold2只開源了模型推理部分的代碼，并沒有開源訓(xùn)練代碼。

這意味著部署好后只能按AlphaFold2設(shè)計(jì)好的流程使用，難以針對(duì)特定的研究項(xiàng)目做一些遷移調(diào)整。

比如拿去預(yù)測(cè)蛋白質(zhì)復(fù)合物結(jié)構(gòu)，研究蛋白質(zhì)與小分子相互作用，又或者與電鏡實(shí)驗(yàn)相結(jié)合，這些都無(wú)能為力。

再加上，該代碼基于谷歌自研的可微分計(jì)算框架JAX及TPU硬件平臺(tái)開發(fā)，高度依賴谷歌生態(tài)系統(tǒng)。

最后，AlphaFold2雖然模型開源，但調(diào)好的參數(shù)組合卻不開放商業(yè)使用。

谷歌母公司Alphabet對(duì)AlphaFold2有著自己的商業(yè)化打算，11月新成立的子公司Isomorphic Laboratories便是其第一步。

鑒于這樣的局限性，像深勢(shì)科技這樣的創(chuàng)業(yè)團(tuán)隊(duì)得自己想辦法。

他們的辦法，便是根據(jù)開源模型自己復(fù)現(xiàn)AlphaFold2的訓(xùn)練部分。

經(jīng)過(guò)幾個(gè)月的攻關(guān)，他們終于在128塊英偉達(dá)V100上復(fù)現(xiàn)了AlphaFold2的全規(guī)模訓(xùn)練。

其中遇到最大的困難是GPU資源有限，難以同時(shí)進(jìn)行多次訓(xùn)練來(lái)比較效果。于是他們只好在設(shè)計(jì)上下功夫，盡量減少試錯(cuò)次數(shù)。

沒想到因此有了一些額外的收獲。

經(jīng)過(guò)混合精度、MPI 并行以及數(shù)據(jù)異步讀取等一系列優(yōu)化后，復(fù)現(xiàn)模型做到了在相同硬件下推理效率還比AlphaFold2公開版高上2-3倍。

另外復(fù)現(xiàn)的模型也不再依賴谷歌生態(tài)，增加了與英偉達(dá)GPU、國(guó)產(chǎn)DPU等硬件體系的適配性。

模型搞定了，那接下來(lái)最關(guān)鍵的問(wèn)題就是，預(yù)測(cè)精度如何？

團(tuán)隊(duì)用CASP14蛋白預(yù)測(cè)任務(wù)中的絕大部分序列進(jìn)行了測(cè)試。

其中包括基于模版建模的兩種難度（TBM-easy/hard），以及最具挑戰(zhàn)性的無(wú)模版建模（FM），發(fā)布時(shí)Uni-Fold平均Cα-lDDT達(dá)到82.6。目前模型還在持續(xù)改進(jìn)中，精度有望再上一層樓。

這個(gè)成績(jī)超過(guò)了RoseTTAFold，與公開版AlphaFold2模型接近。

對(duì)于這個(gè)精度媲美原版、推理效率還更高的蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)工具，他們決定將其推理代碼、訓(xùn)練代碼全部開源，命名為Uni-Fold。

深勢(shì)科技公布訓(xùn)練代碼的舉動(dòng)受到業(yè)界廣泛好評(píng)，連深度學(xué)習(xí)預(yù)測(cè)蛋白質(zhì)結(jié)構(gòu)的先驅(qū)、芝加哥豐田計(jì)算技術(shù)研究所的許錦波教授都評(píng)價(jià)道：

Uni-Fold的出現(xiàn)，讓每個(gè)研究者都可以訓(xùn)練自己的模型。這無(wú)疑將進(jìn)一步推動(dòng)技術(shù)的發(fā)展。對(duì)整個(gè)蛋白質(zhì)折疊領(lǐng)域來(lái)說(shuō)，都是令人激動(dòng)的消息。

深勢(shì)科技這家成立僅3年的公司，為什么能成為國(guó)內(nèi)首家復(fù)現(xiàn)Alphafold2全規(guī)模訓(xùn)練的團(tuán)隊(duì)？

復(fù)現(xiàn)后又為何要率先公開訓(xùn)練代碼？

下面就來(lái)一探究竟。

北大元培系創(chuàng)業(yè)公司，一年完成三輪融資

實(shí)際上，在推出Uni-Fold之前，深勢(shì)科技這個(gè)初創(chuàng)團(tuán)隊(duì)就已經(jīng)頗受業(yè)內(nèi)關(guān)注：

一年狂攬三輪融資，最近一輪達(dá)數(shù)千萬(wàn)美元，由高瓴創(chuàng)投領(lǐng)投。

從團(tuán)隊(duì)背景看，創(chuàng)始人兼首席科學(xué)家張林峰，北京大學(xué)元培學(xué)院出身，在普林斯頓大學(xué)獲應(yīng)用數(shù)學(xué)系博士學(xué)位。

另一位創(chuàng)始人兼CEO孫偉杰同樣來(lái)自北大元培，后繼續(xù)在北京大學(xué)深造，獲管理學(xué)碩士學(xué)位。

孫偉杰在科技領(lǐng)域分析及投資經(jīng)驗(yàn)豐富，張林峰則偏重技術(shù)。他的研究方向包括統(tǒng)計(jì)物理、分子模擬和機(jī)器學(xué)習(xí)及其應(yīng)用。

此外擔(dān)任首席科學(xué)顧問(wèn)的是中科院院士、北京大學(xué)教授鄂維南。

北大元培系A(chǔ)I創(chuàng)業(yè)公司這個(gè)標(biāo)簽，算是深勢(shì)科技受到資本青睞的原因之一。

再看技術(shù)實(shí)力，由鄂維南院士領(lǐng)銜的數(shù)十人科研隊(duì)伍，包含物理建模、數(shù)值算法、機(jī)器學(xué)習(xí)、高性能計(jì)算及藥物和材料計(jì)算等多個(gè)領(lǐng)域人才。

他們大多來(lái)自世界一流高校、科研機(jī)構(gòu)和企業(yè)。

其中許多人擁有交叉背景，深勢(shì)科技特別介紹了一位高中搞過(guò)生物競(jìng)賽，大學(xué)主修人工智能又參加過(guò)超算研究的同學(xué)。

這樣的人了解不同學(xué)科的基礎(chǔ)知識(shí)和術(shù)語(yǔ)體系，對(duì)團(tuán)隊(duì)協(xié)作溝通大有好處。

有著多學(xué)科交叉這樣的特點(diǎn)，深勢(shì)科技自述致力于“以新一代分子模擬方法打造微尺度工業(yè)設(shè)計(jì)平臺(tái)”。

其中的新一代分子模擬方法，指「多尺度建模+機(jī)器學(xué)習(xí)+高性能計(jì)算」新研究范式。

新范式由深勢(shì)科技首創(chuàng)，在保證了計(jì)算效率和精度的基礎(chǔ)上還具備通用性，可為藥物、材料等領(lǐng)域帶來(lái)新的計(jì)算模擬及設(shè)計(jì)工具。

比如鄂維南、張林峰共同參與研究的一項(xiàng)技術(shù)，用機(jī)器學(xué)習(xí)方法將分子動(dòng)力學(xué)極限提升至10億原子規(guī)模同時(shí)保持高精度。

這項(xiàng)成果在2020年獲得有「超算屆諾貝爾獎(jiǎng)」之稱的戈登·貝爾獎(jiǎng)，還與「量子優(yōu)越性」和「人造太陽(yáng)」等成果一起當(dāng)選2020中國(guó)十大科技進(jìn)展。

基于這樣的團(tuán)隊(duì)背景和科研實(shí)力，也就不難理解深勢(shì)科技為何能在這么短的時(shí)間內(nèi)復(fù)現(xiàn)AlphaFold2全規(guī)模訓(xùn)練。

但作為一家創(chuàng)業(yè)公司，光有上述兩種因素還不夠，創(chuàng)造行業(yè)價(jià)值才是當(dāng)前大環(huán)境下的題中之義。

有算法工具的突破做基礎(chǔ)，Unif-Fold配套解決方案也已集成到他們自己的藥物設(shè)計(jì)平臺(tái)Hermite，免費(fèi)注冊(cè)即可測(cè)試使用。

Hermite基于云計(jì)算，無(wú)需操心模型部署、環(huán)境配置問(wèn)題，只要打開瀏覽器就能使用Uni-Fold的各種功能。

除此之外，作為一站式藥物設(shè)計(jì)平臺(tái)的Hermite還集成了更多功能。

如預(yù)測(cè)得到蛋白質(zhì)結(jié)構(gòu)后還可以繼續(xù)進(jìn)行蛋白質(zhì)的動(dòng)力學(xué)模擬、結(jié)構(gòu)精修以及環(huán)區(qū)優(yōu)化。

在藥物優(yōu)化階段也提供了比較不同分子與蛋白靶點(diǎn)結(jié)合活性差別的工具。

這也是深勢(shì)科技打破蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)工具現(xiàn)狀，率先公布訓(xùn)練代碼背后的考量——

他們做的不只是蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)這一個(gè)工具，而是為藥物設(shè)計(jì)領(lǐng)域打造一套模擬-AI-實(shí)驗(yàn)三者驅(qū)動(dòng)的解決方案和開發(fā)平臺(tái)。

而開源開放，吸引更多人來(lái)到開源社區(qū)，所積累的數(shù)據(jù)、匯聚的多學(xué)科人才、實(shí)現(xiàn)的應(yīng)用都是平臺(tái)生態(tài)建設(shè)的重要力量。

深勢(shì)科技的微尺度工業(yè)設(shè)計(jì)平臺(tái)不僅涉及生物醫(yī)藥領(lǐng)域，在新能源材料研發(fā)中也有布局。

最近，深勢(shì)科技與寧德時(shí)代達(dá)成戰(zhàn)略合作，雙方將共建聯(lián)合實(shí)驗(yàn)室，以計(jì)算驅(qū)動(dòng)能源材料設(shè)計(jì)研發(fā)。

計(jì)算生物熱潮

這時(shí)候推出Uni-Fold，倒也并不是心血來(lái)潮，而是聚集了天時(shí)地利人和，選擇謀定而后動(dòng)。

這當(dāng)中既有現(xiàn)實(shí)需求的驅(qū)動(dòng)，也有行業(yè)大勢(shì)的推動(dòng)。

對(duì)深勢(shì)自身而言，過(guò)去一直致力于以AI+物理模型驅(qū)動(dòng)方式為行業(yè)創(chuàng)造價(jià)值。

簡(jiǎn)單來(lái)說(shuō)，就是以第一性原理的方式，將實(shí)際面臨的問(wèn)題或場(chǎng)景演化為定義明確的物理模型。

像蛋白質(zhì)的結(jié)構(gòu)以及構(gòu)象變化、配體與蛋白質(zhì)的相互作用 (包括親和力與動(dòng)力學(xué)參數(shù)的評(píng)估）、以及藥物分子晶型和劑型的溶解度等這些藥物研發(fā)過(guò)程中的關(guān)鍵問(wèn)題，都可以用物理化學(xué)模型進(jìn)行定義。

這些問(wèn)題對(duì)應(yīng)到藥物設(shè)計(jì)平臺(tái)Hermite上，就是諸如藥物優(yōu)化模塊Uni-FEP、蛋白動(dòng)力學(xué)采樣與結(jié)構(gòu)優(yōu)化模塊RiD這類的功能環(huán)節(jié)。

在推出Uni-Fold之前，Hermite就已經(jīng)開啟商業(yè)化進(jìn)程，被多家藥企研發(fā)人員使用。

而Uni-Fold——作為藥物設(shè)計(jì)的第一環(huán)的身份出現(xiàn)，直接為Hermite填補(bǔ)了數(shù)據(jù)驅(qū)動(dòng)引擎。

它有助于深入了解生物機(jī)制，隨后將更多“黑盒子”問(wèn)題轉(zhuǎn)變?yōu)槲锢韱?wèn)題。

中科院院士、北京大學(xué)教授、深勢(shì)科技首席科學(xué)顧問(wèn)鄂維南評(píng)價(jià)道：

物理模型和數(shù)據(jù)驅(qū)動(dòng)結(jié)合的解決方案將為蛋白結(jié)構(gòu)和藥物設(shè)計(jì)行業(yè)提供更好的基礎(chǔ)。

事實(shí)上放眼整個(gè)行業(yè)，計(jì)算生物正在引領(lǐng)一場(chǎng)科技革命。

其核心代表就是以AlphaFold2為代表的新型蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)技術(shù)。

利用原有的實(shí)驗(yàn)手段（比如，冷凍電鏡、X光晶體衍射等）可能需要經(jīng)年累月的研究，才能解析一個(gè)蛋白質(zhì)結(jié)構(gòu)，而利用計(jì)算結(jié)構(gòu)模型最快只需要10分鐘。

而若再向前推進(jìn)，到應(yīng)用端——生物醫(yī)藥領(lǐng)域，新藥研發(fā)的效率有望大幅提高；科學(xué)家能借助AI系統(tǒng)設(shè)計(jì)出自然界不存在的蛋白質(zhì)，催生各種新材料，用于能源、化工、環(huán)保等行業(yè)。

因此，選擇在這時(shí)候完整復(fù)現(xiàn)AlphaFold2，并將其開源、開放給更多的研發(fā)人員，無(wú)疑是最快實(shí)現(xiàn)物盡其用、創(chuàng)造價(jià)值的方式之一。

在人們固有的印象中，以往生命科學(xué)的研究思路都是由表及里、從現(xiàn)象到本質(zhì)。

當(dāng)中也有不少科學(xué)家身先士卒，從本質(zhì)出發(fā)探索微觀分子世界，但始終受限于實(shí)驗(yàn)工具、設(shè)備等外在條件。

直到人工智能的出現(xiàn)，更多自下而上的問(wèn)題逐步得到解決。

算力、算法的發(fā)展能夠真正縮短理論研究與解決實(shí)際問(wèn)題之間的距離。

鄂維南院士曾在多個(gè)場(chǎng)合呼吁：“科學(xué)是人工智能下一個(gè)主戰(zhàn)場(chǎng)。”

AI+Science的范式，價(jià)值已經(jīng)顯現(xiàn)。

而深勢(shì)科技，正是站在這一科技潮頭上的先行者。