新推出的全流程蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)工具M(jìn)EGA-Protein（MindSpore for Evolutionary Generation

& Assessment Protein），其中業(yè)界首創(chuàng) AI MSA 引擎使得蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)突破「孤兒序列」等高精預(yù)測(cè)限制。

日前，昌平實(shí)驗(yàn)室、北京大學(xué)生物醫(yī)學(xué)前沿創(chuàng)新中心（BIOPIC）和化學(xué)與分子工程學(xué)院、深圳灣實(shí)驗(yàn)室高毅勤教授課題組聯(lián)合華為，在中國(guó)算力網(wǎng)智算網(wǎng)絡(luò)關(guān)鍵節(jié)點(diǎn)——西安未來(lái)人工智能計(jì)算中心的昇騰 AI 強(qiáng)大算力支持下，在全場(chǎng)景 AI 框架昇思 MindSpore 上推出全流程蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)工具 MEGA-Protein。[if !vml]

[endif]MEGA-Protein 包含 AI MSA 引擎、蛋白質(zhì)折疊訓(xùn)練推理流程、蛋白質(zhì)結(jié)構(gòu)打分、蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)數(shù)據(jù)集 PSP 等關(guān)鍵技術(shù)，該工具提供了高精度高性能蛋白質(zhì)結(jié)構(gòu)和功能預(yù)測(cè)，其中 AI MSA 引擎能顯著提升單序列的預(yù)測(cè)速度，并且能夠在 MSA 少樣本（few shot）甚至零樣本（zero-shot，即單序列）的情況下，幫助 AlphaFold 2 等模型維持甚至提高推理精度，突破了在「孤兒序列」、高異變序列和人造蛋白等 MSA 匱乏場(chǎng)景下無(wú)法做出準(zhǔn)確預(yù)測(cè)的限制。這是團(tuán)隊(duì)在昇騰 AI 基礎(chǔ)軟硬件平臺(tái)上實(shí)現(xiàn) AlphaFold 2 從訓(xùn)練到推理全流程打通且效率同比提升 2 至 3 倍后，取得的又一次成功。相關(guān)代碼和模型參數(shù)后續(xù)會(huì)在昇思 MindSpore 社區(qū)開(kāi)源，算法和相關(guān)的論文也會(huì)盡快公布。

昇思 MindSpore 首席架構(gòu)師金雪鋒評(píng)論稱：「該工作是產(chǎn)學(xué)界合作的一次重大成功實(shí)踐，也是科學(xué)計(jì)算與人工智能結(jié)合潛力的展示。」金雪鋒補(bǔ)充說(shuō)，支持多計(jì)算范式的 MindSpore 能夠有效提升科研工作的生產(chǎn)力，大幅度提升科學(xué)計(jì)算場(chǎng)景的性能，希望后續(xù)能夠看到更多、更好的成果涌現(xiàn)出來(lái)。高性能蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)引擎 AI MSA，解決「孤兒序列」問(wèn)題蛋白質(zhì)是生命活動(dòng)中最主要的功能體之一，在人體結(jié)構(gòu)和功能中發(fā)揮著核心作用。然而，能夠解析蛋白質(zhì)結(jié)構(gòu)的實(shí)驗(yàn)手段通常都價(jià)格不菲且周期較長(zhǎng)，很難直接通過(guò)它們?nèi)ヌ剿髡麄€(gè)「蛋白質(zhì)宇宙」。蛋白質(zhì)的 3D 結(jié)構(gòu)完全由組成它的氨基酸序列所決定（1972 年諾貝爾獎(jiǎng)得主 Christian Anfinsen 曾通過(guò)實(shí)驗(yàn)證明），這就從理論上確認(rèn)了：存在一個(gè)「理想函數(shù)」，能夠?qū)被嵝蛄?x 映射到其 3D 結(jié)構(gòu) y。因此，蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)問(wèn)題就轉(zhuǎn)化成典型的數(shù)學(xué)建模問(wèn)題。AlphaFold 2 是目前為止人們所能找到的、最接近「理想函數(shù)」的一個(gè)模型。

但是，為了達(dá)到高精度的預(yù)測(cè)，AlphaFold 2 必須引入額外的信息，也即蛋白多序列比對(duì)（MSA）。有多少 MSA 可用，直接決定了 AlphaFold 等模型的預(yù)測(cè)精度。[if !vml]

[endif]圖1. AlphaFold 2 預(yù)測(cè)精度會(huì)隨著可用的 MSA個(gè)數(shù)不足時(shí)而顯著降低。圖片來(lái)源: AlphaFold 2 論文但自然界中仍然存在大量「孤兒序列」，如人類蛋白組中不少蛋白以及大量的病毒等其它抗原相關(guān)蛋白，都存在著 MSA 不夠深的問(wèn)題。此外，在蛋白質(zhì)設(shè)計(jì)領(lǐng)域中，研究人員會(huì)普遍面臨人造序列沒(méi)有 MSA可用的情況。在這些場(chǎng)景下，AlphaFold 2 等模型對(duì)相關(guān)的結(jié)構(gòu)預(yù)測(cè)準(zhǔn)確度會(huì)大幅下降，變得不再適用。

另一方面，AlphaFold 2 等標(biāo)準(zhǔn)檢索 MSA 流程，是在數(shù)據(jù)庫(kù)中使用搜索工具對(duì)目標(biāo)蛋白及相關(guān)序列進(jìn)行多重序列比對(duì)，該流程需要配置數(shù)據(jù)庫(kù)，其中數(shù)據(jù)庫(kù)規(guī)模大概在 1.3T 左右，配置耗時(shí)長(zhǎng)且繁瑣，并且使用數(shù)據(jù)檢索工具（如 HHblits，JackHMMER，MMseqs2 等）搜索時(shí)間長(zhǎng)，不利于科研人員展開(kāi)研究。

昌平實(shí)驗(yàn)室、北京大學(xué)生物醫(yī)學(xué)前沿創(chuàng)新中心（BIOPIC）和化學(xué)與分子工程學(xué)院、深圳灣實(shí)驗(yàn)室高毅勤教授課題組與昇思 MindSpore 社區(qū)新提出的 AI MSA 引擎，對(duì)于原始 MSA 質(zhì)量不高或者數(shù)量少的蛋白，將 AI MSA 引擎接入 AlphaFold 2 后可以明顯改善結(jié)構(gòu)預(yù)測(cè)的質(zhì)量，如下圖所示：[if !vml]

[endif]圖2.?不同方案預(yù)測(cè)的蛋白質(zhì)結(jié)構(gòu)效果對(duì)比圖[if !vml]

[endif]圖3. 零樣本預(yù)測(cè)精度對(duì)比并且端到端的推理性能相較于AlphaFold 2 標(biāo)準(zhǔn)檢索流程有大約 40+ 倍的提升，相較于 MMseqs2 約 5+ 倍的提升，可以大幅提升結(jié)構(gòu)預(yù)測(cè)模型的推理通量，且訓(xùn)練完后的 AI MSA 引擎不需要額外配置數(shù)據(jù)庫(kù)。[if !vml]

[endif]圖4. AI MSA 引擎相關(guān)性能對(duì)比圖團(tuán)隊(duì)表示，AI MSA 引擎是一種對(duì)下游結(jié)構(gòu)模型普適的預(yù)訓(xùn)練方案，可以在免于微調(diào)的情況下直接接入下游的結(jié)構(gòu)預(yù)測(cè)模型，例如 AlphaFold、RoseTTAFold 等，助力科研人員在蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)、蛋白質(zhì)設(shè)計(jì)、蛋白質(zhì)互作等方向進(jìn)行研究，同時(shí)推動(dòng)生物制藥產(chǎn)業(yè)發(fā)展。北京大學(xué)李兆基講席教授謝曉亮認(rèn)為，創(chuàng)新自研的 AI MSA 引擎和基于昇思 MindSpore 的蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)程序的開(kāi)源，標(biāo)志著我們國(guó)家在全球蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)領(lǐng)域名列前茅，該長(zhǎng)足進(jìn)步源于高毅勤團(tuán)隊(duì)在相關(guān)底層技術(shù)上多年的深耕和積累。昇騰 AI 基礎(chǔ)軟硬件平臺(tái)、昇思 MindSpore AI 框架，支持 AI for Science 研究新范式AI MSA 引擎訓(xùn)練參數(shù)量有 47M，數(shù)據(jù)量總共達(dá) 4.4T，其訓(xùn)練對(duì)框架提出了不小的要求，昇思 MindSpore 在昇騰 AI 基礎(chǔ)軟硬件平臺(tái)上與昇騰 CANN 深度結(jié)合，通過(guò)深度協(xié)同優(yōu)化的高性能算子庫(kù)，充分釋放硬件的算力。昇思 MindSpore 采用了多段并行流水線的方式來(lái)構(gòu)建數(shù)據(jù)處理 pipeline，更加細(xì)粒度地規(guī)劃 NPU、CPU 等計(jì)算資源的使用，天然支持各段使用異構(gòu)硬件進(jìn)行流水處理，大大提高了數(shù)據(jù)處理過(guò)程的吞吐量。同時(shí)，昇思 MindSpore 支持大集群高效訓(xùn)練，實(shí)現(xiàn)了優(yōu)質(zhì)的計(jì)算通信比，并且通過(guò)三層 AI 分布式編程范式（手動(dòng)并行＋半自動(dòng)并行＋全自動(dòng)并行），大幅提升分布式并行程序開(kāi)發(fā)效率。相關(guān)代碼和模型參數(shù)后續(xù)會(huì)在昇思 MindSpore 社區(qū)開(kāi)源，算法和相關(guān)的論文也會(huì)盡快公布。團(tuán)隊(duì)還發(fā)布了首個(gè)具有高覆蓋度和多樣性的百萬(wàn)級(jí)蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)數(shù)據(jù)集 PSP。該數(shù)據(jù)集由 570k 個(gè)真實(shí)結(jié)構(gòu)序列（10TB）和 745k 個(gè)互補(bǔ)蒸餾序列（15TB）組成。昇思 MindSpore 基于該數(shù)據(jù)集可進(jìn)行蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)訓(xùn)練，同時(shí)提供了豐富的處理及使用該數(shù)據(jù)集的接口，歡迎試用。