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

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

量子位 報道 | 公眾號 QbitAI

中國版AlphaFold2，來了！

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

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

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

其預(yù)測精度超過RoseTTAFold、與AlphaFold2相當，效率還超AlphaFold2?2-3倍。

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

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

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

這個Uni-Fold究竟有何看頭？接下來一看究竟。

128張GPU煉出中國版AlphaFold2

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

用谷歌官方提供的Colab資源嗎？可以，請排隊。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

團隊用CASP14蛋白預(yù)測任務(wù)中的絕大部分序列進行了測試。

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

這個成績超過了RoseTTAFold，與公開版AlphaFold2模型接近。

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

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

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

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

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

下面就來一探究竟。

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

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

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

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

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

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

此外擔任首席科學顧問的是中科院院士、北京大學教授鄂維南。

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

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

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

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

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

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

其中的新一代分子模擬方法，指「多尺度建模+機器學習+高性能計算」新研究范式。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

計算生物熱潮

這時候推出Uni-Fold，倒也并不是心血來潮，而是聚集了天時地利人和，選擇謀定而后動。

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

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

簡單來說，就是以第一性原理的方式，將實際面臨的問題或場景演化為定義明確的物理模型。

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

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

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

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

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

中科院院士、北京大學教授、深勢科技首席科學顧問鄂維南評價道：

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

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

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

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

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

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

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

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

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

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

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

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

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