Colossal-AI團(tuán)隊聯(lián)合百圖生科開源xTrimo Multimer 推理速度最高提升約 11 倍

　　相比現(xiàn)有方案，Colossal-AI 團(tuán)隊聯(lián)合百圖生科的這項研究，推理速度最高提升約 11 倍!

　　近日，Colossal-AI 團(tuán)隊聯(lián)合百圖生科成功加速蛋白質(zhì)單體與復(fù)合物的結(jié)構(gòu)預(yù)測，行業(yè)內(nèi)最前沿的同時支持蛋白質(zhì)單體(Monomer)與復(fù)合物(Multimer)結(jié)構(gòu)預(yù)測的 xTrimo Multimer 模型已免費(fèi)開源，相比現(xiàn)有方案，推理速度最高提升約 11 倍!

　　本次解決方案背后優(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ī)藥領(lǐng)域落地應(yīng)用的重要實踐，使得在該領(lǐng)域訓(xùn)練和使用更大的模型成為可能。

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

　　Colossal-AI 是一個面向大模型時代的通用深度學(xué)習(xí)系統(tǒng)。自開源以來，Colossal-AI 不僅在數(shù)個月內(nèi)獲得近五千顆 Github Star，多次登上 GitHub Trending 榜首，相關(guān)解決方案也已成功在生物醫(yī)藥、自動駕駛、云計算、零售、芯片等行業(yè)知名廠商落地應(yīng)用，廣受好評。此次攜手百圖生科，利用雙方在高性能計算和生物計算領(lǐng)域上前沿的技術(shù)積淀，提出蛋白質(zhì)單體與復(fù)合物的結(jié)構(gòu)預(yù)測加速方案，進(jìn)一步共同打造繁榮的生物計算生態(tài)，精準(zhǔn)高效地探尋癌癥、衰老等疾病的免疫規(guī)律，逐一攻破 AI 技術(shù)在靶點(diǎn)挖掘、創(chuàng)新藥物設(shè)計落地應(yīng)用上的難關(guān)，讓更多疾病可預(yù)警、可控制、可治愈。

　　眾所周知，蛋白質(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)用中實現(xiàn)超 11 倍的提升。以此為起點(diǎn)，Colossal-AI 團(tuán)隊在蛋白質(zhì)結(jié)構(gòu)預(yù)測領(lǐng)域的技術(shù)革命正在延續(xù)。

　　針對蛋白質(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 市場的剛需。Colossal-AI 致力于構(gòu)建針對 AI 大模型的完整生態(tài)系統(tǒng)，針對不同行業(yè)的需求提供智能化、一體化的高效靈活解決方案，切實有效解決企業(yè) AI 大模型產(chǎn)品落地的關(guān)鍵痛點(diǎn)，為全產(chǎn)業(yè)鏈賦能。