比AlphaFold2快10倍！國(guó)產(chǎn)開(kāi)源項(xiàng)目加速蛋白質(zhì)預(yù)測(cè)，超長(zhǎng)序列22分鐘就搞定

明敏 2022-08-23 13:03:32 來(lái)源：量子位

22分鐘推理出長(zhǎng)達(dá)4K+的超長(zhǎng)蛋白質(zhì)序列，1分15秒搞定將近2K長(zhǎng)的序列預(yù)測(cè)。

這是國(guó)產(chǎn)開(kāi)源項(xiàng)目的最新戰(zhàn)績(jī)！

最近，由Colossal-AI團(tuán)隊(duì)(https://github.com/hpcaitech/ColossalAI)聯(lián)合百圖生科的蛋白質(zhì)預(yù)測(cè)模型xTrimo Multimer，正式免費(fèi)開(kāi)源。

它在支持蛋白質(zhì)單體（Monomer）和復(fù)合物（Multimer）的同時(shí)，還能大幅提升蛋白質(zhì)預(yù)測(cè)速度。

面對(duì)2K到3K序列長(zhǎng)度的預(yù)測(cè)任務(wù)，使用多卡推理，最高速度可達(dá)AlphaFold 2的11.15倍。

還能挑戰(zhàn)4K以上的超長(zhǎng)序列，而AlphaFold 2和OpenFold都受限于顯存，無(wú)法完成這類任務(wù)。

目前，該項(xiàng)目已在GitHub上免費(fèi)開(kāi)源：

https://github.com/hpcaitech/ColossalAI/#xTrimoMultimer

怎么實(shí)現(xiàn)的？

本次解決方案背后優(yōu)化加速技術(shù)來(lái)自Colossal-AI，這是一個(gè)旨在全面助力AI大模型工業(yè)化應(yīng)用的通用深度學(xué)習(xí)系統(tǒng)。

xTrimo Multimer開(kāi)源版模型通過(guò)引入Colossal-AI的大模型優(yōu)化技術(shù)，顯著降低蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)模型訓(xùn)練和推理的時(shí)間和經(jīng)濟(jì)成本，提高蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)模型領(lǐng)域的設(shè)計(jì)和部署效率。

這次Colossal-AI系統(tǒng)在AI大模型&生物醫(yī)藥落地應(yīng)用，使其在該領(lǐng)域訓(xùn)練和使用更大的模型成為可能。

Colossal-AI是一個(gè)面向大模型時(shí)代的通用深度學(xué)習(xí)系統(tǒng)。

自開(kāi)源以來(lái)，它在數(shù)個(gè)月內(nèi)獲得近五千顆Github Star，多次登上GitHub Trending榜首。

相關(guān)解決方案也已成功在生物醫(yī)藥、自動(dòng)駕駛、云計(jì)算、零售、芯片等行業(yè)知名廠商落地應(yīng)用，廣受好評(píng)。

此次和百圖生科聯(lián)手，是結(jié)合了雙方在高性能計(jì)算和生物計(jì)算領(lǐng)域上的前沿技術(shù)積累，提出了蛋白質(zhì)單體與復(fù)合體的結(jié)構(gòu)預(yù)測(cè)加速方案。

這對(duì)于精準(zhǔn)高效探尋癌癥、衰老等疾病的免疫規(guī)律，逐一攻破AI技術(shù)在靶點(diǎn)挖掘、創(chuàng)新藥物設(shè)計(jì)應(yīng)用上，都有所幫助，對(duì)于更多疾病的預(yù)警、控制和治愈都有深遠(yuǎn)意義。

眾所周知，蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)是結(jié)構(gòu)生物學(xué)領(lǐng)域最為重要的課題之一，也是人類理解基因翻譯和蛋白質(zhì)功能的重要手段。

蛋白質(zhì)之間的相互作用，是蛋白質(zhì)發(fā)揮生物學(xué)功能的重要結(jié)構(gòu)基礎(chǔ)。

但是由于蛋白質(zhì)的多級(jí)結(jié)構(gòu)和復(fù)雜的相互作用，使得精確預(yù)測(cè)三維結(jié)構(gòu)這一關(guān)鍵問(wèn)題十分有挑戰(zhàn)性。

近年來(lái)，深度神經(jīng)網(wǎng)絡(luò)的成功使得人工智能在各個(gè)領(lǐng)域得到了廣泛的應(yīng)用。

自DeepMind發(fā)布的AlphaFold能夠?qū)崿F(xiàn)根據(jù)氨基酸序列精準(zhǔn)預(yù)測(cè)蛋白質(zhì)結(jié)構(gòu)以來(lái)，生命科學(xué)領(lǐng)域掀起了使用人工智能預(yù)測(cè)蛋白結(jié)構(gòu)的熱潮。

AlphaFold能夠?qū)崿F(xiàn)端到端的從序列直接預(yù)測(cè)單體蛋白質(zhì)三維結(jié)構(gòu)，但在生物體內(nèi)，多數(shù)蛋白質(zhì)是以復(fù)合體的形式來(lái)發(fā)揮功能的。

因此，為了打破AlphaFold僅在預(yù)測(cè)單體蛋白方面表現(xiàn)優(yōu)異的局限，DeepMind后來(lái)也發(fā)布了用于預(yù)測(cè)蛋白復(fù)合體結(jié)構(gòu)的AlphaFold-Multimer模型。

在2022年3月，Colossal-AI 團(tuán)隊(duì)曾推出蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)模型AlphaFold的訓(xùn)練推理加速方案FastFold，用更低的成本將總體訓(xùn)練時(shí)間從11天減少到67小時(shí)，在長(zhǎng)序列推理應(yīng)用中實(shí)現(xiàn)超11倍的提升。

以此為起點(diǎn)，Colossal-AI 團(tuán)隊(duì)在進(jìn)一步探索蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)領(lǐng)域的技術(shù)革新。

針對(duì)蛋白質(zhì)單體(Monomer)與復(fù)合物(Multimer)結(jié)構(gòu)預(yù)測(cè)的難題，本次Colossal-AI 團(tuán)隊(duì)聯(lián)合百圖生科提出行業(yè)內(nèi)最新解決方案 xTrimo Multimer開(kāi)源版模型，能夠更好地理解蛋白互作關(guān)系，從而提升藥物研發(fā)平臺(tái)中靶點(diǎn)分析、蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)和模擬以及高精準(zhǔn)抗體設(shè)計(jì)等方面的能力。

此外，推理階段昂貴的經(jīng)濟(jì)和時(shí)間成本使得AlphaFold模型相關(guān)的研究和開(kāi)發(fā)十分困難，尤其是面對(duì)長(zhǎng)序列推理時(shí)，計(jì)算復(fù)雜度和內(nèi)存消耗都面臨著巨大挑戰(zhàn)。

xTrimo Multimer開(kāi)源版模型針對(duì)AlphaFold-Multimer模型中的計(jì)算特性，進(jìn)行了針對(duì)性的CUDA優(yōu)化，并且進(jìn)行了Kernel Fusion。

相較于AlphaFold2和哥倫比亞大學(xué)的OpenFold，xTrimo Multimer開(kāi)源版模型在單卡推理上性能有明顯提升 ，推理速度分別提高1.58～2.14倍和1.14~2.23倍。

xTrimo Multimer開(kāi)源版模型還支持超長(zhǎng)序列的分布式推理。

由于使用了動(dòng)態(tài)軸并行（Dynamic Axial Parallelism）技術(shù)，xTrimo Multimer可以高效地將計(jì)算和部分顯存分配到不同的設(shè)備上，從而解決超長(zhǎng)序列面臨的計(jì)算和內(nèi)存挑戰(zhàn)。

在2K到3K的序列長(zhǎng)度下，xTrimo Multimer使用多卡進(jìn)行推理，相對(duì)于OpenFold和AlphaFold 2推理速度最高提升8.47倍和11.15倍，對(duì)比Uni-Fold2.0速度最高提升4.45倍。

xTrimo Multimer可以支持長(zhǎng)達(dá)4K的序列推理，此時(shí)OpenFold和AlphaFold 2受限于顯存無(wú)法完成推理，而xTrimo Multimer可以在20分鐘左右完成。

百圖生科首席AI科學(xué)家宋樂(lè)表示：

從對(duì)單個(gè)蛋白質(zhì)結(jié)構(gòu)的模擬，到對(duì)不同蛋白質(zhì)之間的相互作用的識(shí)別，再到對(duì)蛋白質(zhì)復(fù)合物的繪制，百圖生科致力于解碼、建模復(fù)雜人體免疫系統(tǒng)，開(kāi)發(fā)突破創(chuàng)新藥物，編程免疫系統(tǒng)，治愈多種免疫相關(guān)疾病。

此次xTrimo Multimer開(kāi)源版模型的發(fā)布，是我們攜手潞晨科技Colossal-AI 團(tuán)隊(duì)，借助其在高性能計(jì)算上的優(yōu)勢(shì)以及百圖生科生物計(jì)算領(lǐng)域前沿的技術(shù)積淀，朝著百圖生科xTrimo多模態(tài)生物計(jì)算大模型體系邁出的又一步。

正如我們一直以來(lái)所努力的方向，精準(zhǔn)高效地探尋癌癥、衰老等疾病的免疫規(guī)律，逐一攻破AI技術(shù)在靶點(diǎn)挖掘、創(chuàng)新藥物設(shè)計(jì)落地應(yīng)用上的難關(guān)，讓更多疾病可預(yù)警、可控制、可治愈！

潞晨科技創(chuàng)始人尤洋教授表示：

此次 Colossal-AI 團(tuán)隊(duì)與百圖生科合作的蛋白質(zhì)單體與復(fù)合物結(jié)構(gòu)預(yù)測(cè)最新方案，是面向大模型時(shí)代的通用深度學(xué)習(xí)系統(tǒng) Colossal-AI 在生物醫(yī)藥領(lǐng)域應(yīng)用落地的重要進(jìn)展。

未來(lái)我們雙方將會(huì)繼續(xù)在生物計(jì)算大模型上有更深入的合作，助力深度學(xué)習(xí)在創(chuàng)新藥物研發(fā)上的應(yīng)用和落地。

隨著人工智能對(duì)各個(gè)領(lǐng)域的高度滲透，各行業(yè)的智能化產(chǎn)業(yè)轉(zhuǎn)型升級(jí)都離不開(kāi)高性能AI的加持，而潞晨的AI大模型的高效訓(xùn)練和推理方案，剛好可為AI市場(chǎng)提供解決方案。

團(tuán)隊(duì)介紹

關(guān)于百圖生科

百圖生科(BioMap)是生物計(jì)算引擎驅(qū)動(dòng)的突破創(chuàng)新藥物研發(fā)平臺(tái),由百度創(chuàng)始人李彥宏發(fā)起創(chuàng)立,致力于將先進(jìn)AI等信息技術(shù)(IT)與前沿生物技術(shù)(BT)相結(jié)合,構(gòu)建獨(dú)特的靶點(diǎn)挖掘及藥物設(shè)計(jì)能力,聚焦于解析免疫調(diào)控機(jī)理,開(kāi)發(fā)創(chuàng)新的治療性藥物,造福人類健康。

目前公司擁有近萬(wàn)平米的實(shí)驗(yàn)室，50+個(gè)靶點(diǎn)及藥物研發(fā)資產(chǎn)組合。

關(guān)于潞晨科技

潞晨科技致力于解放AI生產(chǎn)力，打造面向大模型時(shí)代的通用深度學(xué)習(xí)系統(tǒng)Colossal-AI ，高效促進(jìn)AI大模型落地應(yīng)用。

核心成員來(lái)自美國(guó)加州伯克利、斯坦福、清華、北大、新加坡國(guó)立、南洋理工大學(xué)等世界一流高校，在國(guó)際頂級(jí)學(xué)術(shù)刊物或會(huì)議共發(fā)表論文近百篇，曾在谷歌、微軟、NVIDIA、IBM、英特爾等頭部科技公司任職。

公司成立即獲得位列“清科中國(guó)早期投資機(jī)構(gòu)30強(qiáng)”榜首的創(chuàng)新工場(chǎng)、真格基金超千萬(wàn)元種子輪投資。

傳送門

GitHub地址：

https://github.com/hpcaitech/ColossalAI/#xTrimoMultimer
https://github.com/hpcaitech/ColossalAI

參考鏈接：
https://www.technologyreview.com/2021/07/22/1029973/deepmind-alphafold-protein-folding-biology-disease-drugs-proteome/