入選ICML 2025,清華/人大/字節(jié)提出首個(gè)跨分子種類統(tǒng)一生成框架UniMoMo,實(shí)現(xiàn)多類型藥物分子設(shè)計(jì) 原創(chuàng)
清華大學(xué)劉洋老師組、人民大學(xué)高瓴人工智能學(xué)院黃文炳老師組、字節(jié)跳動(dòng) AI 制藥團(tuán)隊(duì)共同提出了一種跨分子種類統(tǒng)一生成框架 UniMoMo 。該框架基于分子片段(block)對(duì)不同種類的分子進(jìn)行統(tǒng)一表示,使用變分自編碼器對(duì)每個(gè) block 的全原子構(gòu)象進(jìn)行壓縮,并在壓縮后的隱空間進(jìn)行幾何擴(kuò)散建模(diffusion),從而實(shí)現(xiàn)對(duì)同一靶點(diǎn)不同結(jié)合分子種類(小分子、多肽、抗體)的設(shè)計(jì)。 UniMoMo 在多類分子任務(wù)基準(zhǔn)的評(píng)測(cè)中均實(shí)現(xiàn)領(lǐng)先表現(xiàn),展示了跨模態(tài)知識(shí)遷移與數(shù)據(jù)共享的巨大潛力。
相關(guān)成果以「UniMoMo: Unified Generative Modeling of 3D Molecules for De Novo Binder Design」為題,入選 ICML 2025 。
論文地址:
???https://go.hyper.ai/wZXZZ????
開源項(xiàng)目地址:
??https://github.com/kxz18/UniMoMo??????
為什么需要統(tǒng)一建模
不同的分子類型在藥物開發(fā)中各有優(yōu)劣,因此在不同的疾病場(chǎng)景中往往需要選擇最合適的分子類型。例如:
* 小分子體積小、易于口服、穿透性強(qiáng),適合進(jìn)入細(xì)胞內(nèi)部作用于靶點(diǎn),廣泛應(yīng)用于慢性病和代謝類疾病等;
* 多肽類分子擁有較高的靶向性,能夠與蛋白質(zhì)表面的大而平的區(qū)域結(jié)合,適合靶向「難成藥」的蛋白互作位點(diǎn),常用于癌癥、炎癥等治療中;
* 抗體具有極高的選擇性和親和力,能夠穩(wěn)定識(shí)別特定的蛋白標(biāo)志物,特別適合免疫治療等精確干預(yù)場(chǎng)景。
因此,面對(duì)不同的疾病機(jī)制、靶點(diǎn)特征和用藥需求,適合使用的分子類型是不同的?,F(xiàn)有的生成方法通常只針對(duì)某一類分子(如小分子、肽或抗體)進(jìn)行建模,
既無(wú)法滿足多樣化的治療需求,也無(wú)法利用不同分子之間的共性來(lái)提升模型性能。
從應(yīng)用角度來(lái)看,統(tǒng)一建??梢宰屛覀冡槍?duì)同一個(gè)靶點(diǎn)同時(shí)探索多種類型的候選藥物,為不同的下游場(chǎng)景提供更多選擇。
從機(jī)器學(xué)習(xí)角度出發(fā),不同分子類型之間共享相似的結(jié)合規(guī)律(氫鍵、π-π 堆疊、鹽橋等)和幾何約束(鍵長(zhǎng)、鍵角等),是可以相互借鑒的,
因此統(tǒng)一建模通過利用更大的數(shù)據(jù)規(guī)模,應(yīng)當(dāng)能提升模型的泛化能力和交叉遷移能力。
UniMoMo 為同一結(jié)合位點(diǎn)設(shè)計(jì)不同種類的結(jié)合分子
生成式統(tǒng)一建模的困難
盡管統(tǒng)一生成不同類型分子的想法令人期待,但要實(shí)現(xiàn)這樣一個(gè)框架仍面臨巨大挑戰(zhàn),主要體現(xiàn)在分子表示的選擇,以及生成算法的設(shè)計(jì)上。
首先,不同分子類型在結(jié)構(gòu)表示上存在較大差異:小分子由各種官能團(tuán)組成,結(jié)構(gòu)高度多樣且非線性;而肽和抗體則由氨基酸以線性序列連接構(gòu)成,尤其抗體還具有功能明確的功能區(qū)域劃分。一種直觀但效果不佳的做法是將所有分子都視為原子圖來(lái)建模。
然而這種方式忽略了分子天然的層次結(jié)構(gòu),比如苯環(huán)或標(biāo)準(zhǔn)氨基酸等關(guān)鍵子結(jié)構(gòu),并且導(dǎo)致在處理如抗體這類結(jié)合面較大的體系時(shí)計(jì)算成本極高。
反之,如果僅使用常見的結(jié)構(gòu)片段詞表構(gòu)建立片段級(jí)別的圖(例如多數(shù)蛋白質(zhì)設(shè)計(jì)的工作,只考慮 Cα 坐標(biāo)),
而忽略原子級(jí)細(xì)節(jié),則會(huì)犧牲分子生成的可遷移性與精度,因?yàn)榻Y(jié)合分子設(shè)計(jì)的本質(zhì)規(guī)律是與靶點(diǎn)的空間相互作用,以及分子內(nèi)部的幾何約束,這些都是定義在原子級(jí)別上的物理規(guī)律,需要精確的全原子信息支持。
因此,要構(gòu)建一個(gè)真正有效且高效的統(tǒng)一分子表示,就必須同時(shí)解決兩個(gè)難題:
既要保留原子層面的幾何細(xì)節(jié),又要抽象出結(jié)構(gòu)上的層級(jí)先驗(yàn)。
其次,如果在生成中引入結(jié)構(gòu)片段來(lái)保留層級(jí)先驗(yàn),則對(duì)生成算法帶來(lái)核心挑戰(zhàn):
傳統(tǒng)的擴(kuò)散模型通常依賴于定長(zhǎng)、固定結(jié)構(gòu)的數(shù)據(jù)表示,例如固定數(shù)量的點(diǎn)云或原子。對(duì)于 AF3 一類的結(jié)構(gòu)預(yù)測(cè)模型而言,由于 2D 拓?fù)涫翘崆敖o定的,擴(kuò)散過程中并不會(huì)產(chǎn)生原子數(shù)或 2D 結(jié)構(gòu)的變化。而對(duì)于分子生成的任務(wù),2D 拓?fù)浜?3D 結(jié)構(gòu)是需要被同時(shí)生成的,而當(dāng)去噪過程中結(jié)構(gòu)片段的類型發(fā)生變化時(shí),其所對(duì)應(yīng)的原子數(shù)量、種類和排布都會(huì)隨之改變。這一點(diǎn)打破了常規(guī)擴(kuò)散模型的假設(shè),對(duì)建模提出了極高要求。
UniMoMo:統(tǒng)一的生成模型
為了解決不同分子類型結(jié)構(gòu)差異大、建模難度高的問題,文章提出了一個(gè)全新的框架——UniMoMo,
它從兩個(gè)關(guān)鍵設(shè)計(jì)入手,有效兼顧了結(jié)構(gòu)的層次性和原子級(jí)精度:
* 統(tǒng)一表示:以 block 形式建模所有分子類型。
無(wú)論是小分子、肽還是抗體,UniMoMo 都將其結(jié)構(gòu)表示為由分子片段(block)組成的圖。其中每個(gè) block 可以是一個(gè)標(biāo)準(zhǔn)氨基酸,也可以是一個(gè)常見的小分子片段(如苯環(huán)、吲哚等)。在文章的實(shí)現(xiàn)中,記錄的分子片段包括所有標(biāo)準(zhǔn)氨基酸和由主子圖(principle subgraph)挖掘算法自動(dòng)識(shí)別而得的小分子片段。所有的非天然氨基酸可以當(dāng)作小分子進(jìn)行分詞。
這種表示方式既同時(shí)保留了分子的原子級(jí)別細(xì)節(jié),以及不同種類分子自身的層級(jí)結(jié)構(gòu),使統(tǒng)一建模成為可能。
* 全原子幾何隱空間擴(kuò)散模型:在壓縮表示上進(jìn)行高效生成。
為了解決生成過程中 block 種類改變導(dǎo)致的原子種類和數(shù)量同步改變的問題,并且為了提升生成效率和結(jié)構(gòu)精度,
文章設(shè)計(jì)了一個(gè)全原子的迭代變分自編碼器(IterVAE),將每個(gè) block 中的所有原子壓縮為隱空間中的一個(gè)「點(diǎn)」,包括定長(zhǎng)的隱空間表示向量和對(duì)應(yīng)的隱空間坐標(biāo)。
然后,模型在這個(gè)壓縮的幾何隱空間中進(jìn)行生成式建模,生成新分子的潛在表示,最后再解碼回完整的原子結(jié)構(gòu)。
由于隱空間的數(shù)據(jù)表示是定長(zhǎng)(預(yù)先給定 block 的數(shù)量)且連續(xù)的,因此可以很容易地兼容已有的各種生成算法。在目前的嘗試中,擴(kuò)散模型(diffusion)已經(jīng)能產(chǎn)生比較不錯(cuò)的結(jié)果。這種設(shè)計(jì)使得模型在生成過程中聚焦于 block 之間的全局布局,而細(xì)致的原子級(jí)構(gòu)造則由解碼器來(lái)補(bǔ)全,從而同時(shí)實(shí)現(xiàn)了高效性與原子級(jí)準(zhǔn)確性的統(tǒng)一。
不同分子種類的分詞方案以及模型整體架構(gòu)
統(tǒng)一建模超越單域建模
為了驗(yàn)證 UniMoMo 在不同分子類型上的通用性與有效性,作者在多個(gè)基于結(jié)構(gòu)設(shè)計(jì)任務(wù)中進(jìn)行了系統(tǒng)評(píng)估,
涵蓋了小分子、多肽和抗體三類代表性的結(jié)合分子。通過與對(duì)應(yīng)領(lǐng)域內(nèi)最具代表性的單分子類型生成模型進(jìn)行對(duì)比,實(shí)驗(yàn)旨在探討統(tǒng)一建模是否具備更強(qiáng)的幾何建模能力與跨模態(tài)泛化能力,特別是在空間結(jié)構(gòu)合理性與結(jié)合能力等關(guān)鍵指標(biāo)上的表現(xiàn)。
結(jié)果顯示,
統(tǒng)一訓(xùn)練的 UniMoMo 在所有分子類型上都實(shí)現(xiàn)了全面超越,不僅在結(jié)構(gòu)還原精度上表現(xiàn)優(yōu)異,更在關(guān)鍵的幾何合理性和與靶點(diǎn)的相互作用質(zhì)量上取得顯著提升。
與靶點(diǎn)結(jié)合的多肽設(shè)計(jì)結(jié)果
在多肽生成任務(wù)中,UniMoMo 在多個(gè)關(guān)鍵指標(biāo)上顯著優(yōu)于現(xiàn)有的專域模型,包括 RFDiffusion 、 PepFlow 和 PepGLAD 等。尤其在結(jié)構(gòu)準(zhǔn)確性方面,UniMoMo 達(dá)到了更低的復(fù)合物 RMSD 和單體的 RMSD,表明其生成的肽結(jié)構(gòu)更加接近真實(shí)結(jié)合構(gòu)象。
UniMoMo 還能生成具有更低 Rosetta 結(jié)合能的結(jié)構(gòu),
反映出其對(duì)蛋白結(jié)合位點(diǎn)幾何特征的更強(qiáng)建模能力。此外,在衡量多肽構(gòu)象質(zhì)量的二面角度分布一致性(JSD of backbone/sidechain torsions)以及原子級(jí)空間沖突(clash rate)等幾何合理性指標(biāo)上,UniMoMo 同樣展現(xiàn)出領(lǐng)先表現(xiàn)。并且,使用所有數(shù)據(jù)訓(xùn)練的 UniMoMo(all)在各種指標(biāo)上一致性地超過了僅使用多肽數(shù)據(jù)訓(xùn)練的模型,
證明了 UniMoMo 的跨分子種類學(xué)習(xí)和泛化的能力。
抗體 CDR-H3 的設(shè)計(jì)結(jié)果
抗體 CDR-H3 的設(shè)計(jì)的合理性結(jié)果
在抗體設(shè)計(jì)任務(wù)中,UniMoMo 同樣展現(xiàn)出強(qiáng)勁性能。相比現(xiàn)有方法如 MEAN 、 dyMEAN 和 DiffAb,
UniMoMo 在召回天然結(jié)合的序列和結(jié)構(gòu)(AAR 和 RMSD)以及結(jié)合能提升(IMP)等關(guān)鍵指標(biāo)上均實(shí)現(xiàn)全面超越。特別是在多次采樣生成的評(píng)估中,UniMoMo 能夠以更高概率生成接近天然構(gòu)象的抗體片段,顯示出其在抗體結(jié)構(gòu)空間中的良好探索能力。
同樣的,使用跨分子類型數(shù)據(jù)聯(lián)合訓(xùn)練的 UniMoMo(all)在所有指標(biāo)上均優(yōu)于僅使用抗體數(shù)據(jù)訓(xùn)練的版本,
說明統(tǒng)一建模確實(shí)有助于模型學(xué)習(xí)到更普適、更可遷移的分子結(jié)構(gòu)空間規(guī)律。這一結(jié)果凸顯了不同分子類型之間在結(jié)構(gòu)建模上的共性,也驗(yàn)證了跨域數(shù)據(jù)融合對(duì)提升生成質(zhì)量的顯著價(jià)值。
小分子設(shè)計(jì)的整體比較
小分子設(shè)計(jì)的子結(jié)構(gòu)分析
在小分子生成任務(wù)中,UniMoMo 同樣展現(xiàn)了優(yōu)越的性能。通過在 CrossDocked2020 數(shù)據(jù)集上的評(píng)估,
作者發(fā)現(xiàn) UniMoMo 在基于 CBGBench 的綜合評(píng)價(jià)上超越了現(xiàn)有主流方法。
具體而言,UniMoMo 在子結(jié)構(gòu)分布(原子種類、官能團(tuán)等)、化學(xué)性質(zhì)合理性(QED 、 LogP 、 SA 等)、幾何結(jié)構(gòu)質(zhì)量(鍵長(zhǎng)/角度分布與原子沖突率等)以及相互作用得分(Vina docking)等方面均取得了更高的綜合評(píng)分(完整的實(shí)驗(yàn)結(jié)果請(qǐng)查看原文)。尤為重要的是,與僅在小分子數(shù)據(jù)上訓(xùn)練的單域版本相比,跨分子類型訓(xùn)練的 UniMoMo(all)在所有評(píng)估維度上均有顯著提升。這表明,
即便是在分子結(jié)構(gòu)構(gòu)成最為靈活、類型最為多樣的小分子場(chǎng)景中,統(tǒng)一模型依然能夠從其他分子類型中遷移幾何規(guī)律與相互作用模式,進(jìn)而提升小分子的單體構(gòu)象和相對(duì)口袋空間布局的合理性。這一現(xiàn)象再次驗(yàn)證了 UniMoMo 的核心理念:不同分子之間的幾何約束和結(jié)合機(jī)制具有可共享的模式,統(tǒng)一建模可以有效激發(fā)這一潛力。
綜合三類任務(wù)的實(shí)驗(yàn)結(jié)果,UniMoMo 展現(xiàn)出高度一致的優(yōu)勢(shì):使用跨分子種類數(shù)據(jù)訓(xùn)練的統(tǒng)一模型在各自任務(wù)中優(yōu)于現(xiàn)有的單域生成模型,且相比只用單域數(shù)據(jù)訓(xùn)練的 UniMoMo 有明顯的能力提升。這一現(xiàn)象表明,分子設(shè)計(jì)中看似截然不同的任務(wù),其底層的物理化學(xué)約束與空間幾何規(guī)律實(shí)則具有高度共性,
而 UniMoMo 的統(tǒng)一建模策略正好捕捉并放大了這種共性,從而實(shí)現(xiàn)跨任務(wù)遷移、互補(bǔ)增強(qiáng)。這些結(jié)果不僅驗(yàn)證了 UniMoMo 的有效性,也為未來(lái)構(gòu)建更強(qiáng)大的統(tǒng)一的分子生成系統(tǒng)提供了強(qiáng)有力的實(shí)證支持。
GPCR 案例研究
UniMoMo 在 GPCR 相同結(jié)合位點(diǎn)設(shè)計(jì)的不同結(jié)合分子
作為案例,作者選取了人類最重要的藥物靶標(biāo)之一——G 蛋白偶聯(lián)受體(GPCR),評(píng)估 UniMoMo 在同一結(jié)合位點(diǎn)上生成不同類型分子(多肽、抗體、小分子)的能力。
UniMoMo 所生成的多肽、抗體和小分子常用結(jié)合能評(píng)價(jià)的力場(chǎng)(Rosetta ΔG,Vina score)下均展現(xiàn)出不錯(cuò)的分布。而更令人驚喜的是,生成的小分子結(jié)構(gòu)中還自發(fā)地模擬出類似天然氨基酸側(cè)鏈的官能團(tuán),用于構(gòu)建氫鍵,與靶點(diǎn)形成關(guān)鍵互作。此外,小分子還借鑒了多肽與抗體中的局部幾何構(gòu)型,如分子骨架上的酰胺連接,使其能夠有效填充原本更適合大分子的結(jié)合口袋。這一案例生動(dòng)展示了 UniMoMo 在實(shí)際任務(wù)中跨模態(tài)借鑒、自動(dòng)適應(yīng)結(jié)合口袋的能力,體現(xiàn)了其在三維結(jié)構(gòu)層面深度理解靶點(diǎn)與分子之間相互作用,以及分子內(nèi)部幾何約束的潛力。
未來(lái)的探索
盡管 UniMoMo 已在多個(gè)分子類型與任務(wù)中展現(xiàn)出強(qiáng)大的統(tǒng)一生成能力,但作者也指出,該方向仍存在諸多值得探索的未來(lái)可能。
當(dāng)前工作主要聚焦于天然氨基酸與常見分子片段的建模,后續(xù)可進(jìn)一步擴(kuò)展到非天然氨基酸、后修飾多肽/抗體、環(huán)狀分子等結(jié)構(gòu)更加復(fù)雜的藥物形式,從而覆蓋更廣泛的候選分子空間。統(tǒng)一建模的理念也為模型的可控性和可解釋性研究提供了契機(jī),有望進(jìn)一步推進(jìn)生成模型向更可信、更實(shí)用的分子設(shè)計(jì)平臺(tái)發(fā)展??傊?,UniMoMo 的提出不僅為分子設(shè)計(jì)任務(wù)提供了一個(gè)通用而強(qiáng)大的生成框架,也為 AI 驅(qū)動(dòng)的藥物發(fā)現(xiàn)打開了一個(gè)充滿潛力的新方向。
