比AlphaFold2快10倍！國產(chǎn)開源項目加速蛋白質(zhì)預(yù)測，超長序列22分鐘就搞定

明敏 2022-08-23 13:03:32 來源：量子位

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

這是國產(chǎn)開源項目的最新戰(zhàn)績！

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

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

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

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

目前，該項目已在GitHub上免費(fèi)開源：

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

xTrimo Multimer開源版模型還支持超長序列的分布式推理。

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

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

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

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

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

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

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

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

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

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

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

團(tuán)隊介紹

關(guān)于百圖生科

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

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

關(guān)于潞晨科技

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

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

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

傳送門

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/