12月8日，來自深勢科技的團(tuán)隊推出了蛋白質(zhì)結(jié)構(gòu)預(yù)測工具Uni-Fold，并開源訓(xùn)練、推理代碼。Uni-Fold成功復(fù)現(xiàn)了曾引起生物學(xué)界轟動的Alphafold2的全規(guī)模訓(xùn)練，相應(yīng)解決方案已集成至深勢科技自主研發(fā)的藥物設(shè)計平臺Hermite，供廣大用戶測試使用。Uni-Fold克服了Alphafold2未開源訓(xùn)練代碼、硬件支持單一、模型不可商用等局限性，在訓(xùn)練與推理環(huán)節(jié)進(jìn)行了GPU上的適配、性能優(yōu)化、功能完善等工作，為更多人參與推動領(lǐng)域進(jìn)一步發(fā)展提供了基礎(chǔ)。
圖1 Uni-Fold預(yù)測的蛋白結(jié)構(gòu)
蛋白質(zhì)是組成生命體的重要物質(zhì)基礎(chǔ)。不同的蛋白質(zhì)通過折疊形成不同的三維結(jié)構(gòu)，執(zhí)行多種多樣的生理功能。從氨基酸序列到對應(yīng)蛋白質(zhì)三維結(jié)構(gòu)的預(yù)測問題被認(rèn)為是生物學(xué)領(lǐng)域最具有挑戰(zhàn)性的問題之一。在過去幾十年的研究中，科學(xué)家們?yōu)樘綔y蛋白質(zhì)結(jié)構(gòu)發(fā)展出三大實驗手段：X射線晶體學(xué)、核磁共振和冷凍電鏡。但實驗方法成本高、周期長。目前人類已知有數(shù)十億氨基酸序列，但還原出結(jié)構(gòu)的只有十幾萬。生物學(xué)發(fā)展因此頗受掣肘。
近年來AI的發(fā)展為這一問題帶來了新的可能。在2020年的國際蛋白質(zhì)結(jié)構(gòu)預(yù)測頂級競賽CASP14Critical Assessment of protein Structure Prediction 14上，Google DeepMind推出的AlphaFold2預(yù)測的大部分結(jié)構(gòu)準(zhǔn)確度均與實驗足夠接近，遠(yuǎn)超此前所有方法。這一結(jié)果也標(biāo)志著蛋白質(zhì)單體的結(jié)構(gòu)問題在一定意義上得到了解決。
不過，盡管DeepMind在半年后開源了AlphaFold2模型的推理代碼，其訓(xùn)練代碼并未被公布，且模型不可商用。研究者難以通過訓(xùn)練或調(diào)整模型進(jìn)一步提升AlphaFold2的表現(xiàn)，或遷移至更多應(yīng)用場景，如蛋白質(zhì)復(fù)合物結(jié)構(gòu)預(yù)測、蛋白質(zhì)與小分子相互作用、與電鏡實驗結(jié)合等。此外，該代碼基于Google自研的可微分計算框架JAX及TPU硬件平臺開發(fā)，高度依賴Google生態(tài)系統(tǒng)。
Uni-Fold打破了這些限制。目前已開源的Uni-Fold訓(xùn)練代碼，經(jīng)過MPI并行、數(shù)據(jù)異步讀取、混合精度訓(xùn)練等優(yōu)化工作，支持高效GPU訓(xùn)練。基于優(yōu)化后的訓(xùn)練代碼，深勢科技團(tuán)隊實現(xiàn)在128塊NVIDIA GPU上進(jìn)行AlphaFold2模型的全規(guī)模訓(xùn)練，并將持續(xù)優(yōu)化所得模型。Uni-Fold的推理代碼更加輕量、高效，在相同硬件環(huán)境下，可獲得與公布的Alphafold2代碼相比2-3倍的效率提升。
在相同的測試條件下，Uni-Fold的預(yù)測精度超越了華盛頓大學(xué)蛋白設(shè)計研究所David Baker教授課題組研發(fā)的RoseTTAFold，與DeepMind官方發(fā)布的AlphaFold2模型接近。深勢科技團(tuán)隊使用Uni-Fold對CASP14蛋白預(yù)測任務(wù)中的序列數(shù)據(jù)集進(jìn)行了直接測試，平均Cα-lDDT達(dá)到82.6。可以說，Uni-Fold是“訓(xùn)練公開版的AlphaFold2”，而且硬件適配與效率更加優(yōu)化。
中科院院士、北京大學(xué)教授、深勢科技首席科學(xué)顧問鄂維南評價道：“盡管DeepMind開源了推理代碼，但模型的訓(xùn)練技術(shù)才是核心競爭力。在資源比較匱乏的情況下，深勢科技團(tuán)隊在較短時間內(nèi)實現(xiàn)模型訓(xùn)練、開源訓(xùn)練代碼，這是很不容易的。深勢科技在物理模型方面具備豐富的經(jīng)驗，Uni-Fold進(jìn)一步彌補(bǔ)了團(tuán)隊在數(shù)據(jù)驅(qū)動方面的短板。下一步，物理模型和數(shù)據(jù)驅(qū)動結(jié)合的解決方案將為蛋白結(jié)構(gòu)和藥物設(shè)計行業(yè)提供更好的基礎(chǔ)。”
深勢科技也對Uni-Fold進(jìn)行了進(jìn)一步的開發(fā)和優(yōu)化，集成于旗下自主研發(fā)的一站式計算輔助藥物設(shè)計平臺Hermite中，供用戶使用。Hermite是根植于云計算的臨床前藥物設(shè)計全流程計算平臺，致力于為藥物科研工作者提供數(shù)據(jù)、算法、算力三位一體的一站式解決方案。在Uni-Fold前，Hermite上已集成多個重量級功能。例如基于Uni-Fold等方法得到蛋白結(jié)構(gòu)后進(jìn)行蛋白的動力學(xué)模擬、結(jié)構(gòu)精修或環(huán)區(qū)優(yōu)化的RiD模塊;在藥物優(yōu)化階段比較不同分子與蛋白靶點結(jié)合活性差別的Uni-FEP模塊，等等。
圖2 Uni-Fold在深勢科技Hermite藥物設(shè)計平臺上的使用
在藥物設(shè)計領(lǐng)域，結(jié)合Uni-Fold、RiD等AI輔助計算模擬的算法，以及算法與冷凍電鏡等實驗手段的結(jié)合，深勢科技正在打造出一套模擬―AI―實驗三者驅(qū)動的解決方案，賦能藥物設(shè)計行業(yè)，加速從業(yè)者的藥物研發(fā)進(jìn)程。
芝加哥豐田計算技術(shù)研究所的許錦波教授主攻機(jī)器學(xué)習(xí)、優(yōu)化方法和計算生物學(xué)，也是深度學(xué)習(xí)預(yù)測蛋白質(zhì)結(jié)構(gòu)的先驅(qū)。他評價道：“針對AlphaFold未公開訓(xùn)練代碼的局限性，深勢科技敢于打破現(xiàn)狀，率先公布訓(xùn)練代碼。Uni-Fold的出現(xiàn)，讓每個研究者都可以訓(xùn)練自己的模型。這無疑將進(jìn)一步推動技術(shù)的發(fā)展。對整個蛋白質(zhì)折疊領(lǐng)域來說，都是令人激動的消息。”
