入選NeurIPS 2024！西湖大學(xué)提出通用分子逆折疊模型UniIF，對(duì)AlphaFold 3形成

分子逆折疊在藥物和材料設(shè)計(jì)中起到關(guān)鍵作用，使得科學(xué)家能夠合成具有理想結(jié)構(gòu)的新分子。過去的研究大多集中于大分子或小分子的逆折疊，但卻很少關(guān)注通用分子的逆折疊。

**構(gòu)建統(tǒng)一的通用模型主要存在 3 大挑戰(zhàn)：**① 單位差異：大分子一般將預(yù)定義的微結(jié)構(gòu)作為基本單元，如氨基酸之于蛋白質(zhì)，核苷酸之于 RNA；而小分子將原子作為基本單元；② 幾何特征提?。翰煌芯吭趲缀翁卣魈崛》矫娌捎枚喾N策略，如距離、角度和張量積，缺乏統(tǒng)一的特征化方法；③ 系統(tǒng)規(guī)模：小分子允許全局注意力機(jī)制來學(xué)習(xí)長(zhǎng)期依賴關(guān)系，但這往往在大分子上行不通。

為了解決上述挑戰(zhàn)，并進(jìn)一步補(bǔ)充 RoseTTAFold All-Atom 和 AlphaFold 3 在分子結(jié)構(gòu)預(yù)測(cè)上取得的進(jìn)展，**來自西湖大學(xué)未來產(chǎn)業(yè)研究中心的團(tuán)隊(duì)提出了一個(gè)統(tǒng)一模型 UniIF，用于所有分子的逆折疊。**研究人員在蛋白質(zhì)設(shè)計(jì)、RNA 設(shè)計(jì)和材料設(shè)計(jì)等多個(gè)任務(wù)上進(jìn)行了全面實(shí)驗(yàn)，以證明 UniIF 的有效性。結(jié)果表明，UniIF 在所有任務(wù)上都達(dá)到了最先進(jìn)的性能。

相關(guān)研究以「UniIF: Unified Molecule Inverse Folding」為題，入選頂會(huì) NeurIPS 2024。

研究亮點(diǎn)：

* 研究提出的統(tǒng)一模型 UniIF 為一般分子逆折疊提供了一種多功能且有效的解決方案

* 該模型從兩個(gè)層面進(jìn)行統(tǒng)一：在數(shù)據(jù)層面提出了所有分子的統(tǒng)一塊圖數(shù)據(jù)形式，包括局部坐標(biāo)系的構(gòu)建和幾何特征的初始化；在模型層面引入幾何塊注意力網(wǎng)絡(luò)，捕捉所有分子的三維相互作用

* 研究人員證明了所提出的方法在蛋白質(zhì)設(shè)計(jì)、RNA 設(shè)計(jì)和材料設(shè)計(jì)三大任務(wù)上都優(yōu)于最先進(jìn)的方法，這一成就可能對(duì)機(jī)器學(xué)習(xí)、藥物發(fā)現(xiàn)和材料科學(xué)界產(chǎn)生積極影響

**數(shù)據(jù)集：選擇對(duì)應(yīng)數(shù)據(jù)集進(jìn)行三種任務(wù)實(shí)驗(yàn)

蛋白質(zhì)設(shè)計(jì)任務(wù)中**，研究人員在 CATH4.3 數(shù)據(jù)集上評(píng)估 UniIF。該數(shù)據(jù)集按 CATH 拓?fù)浞诸惔a分割，產(chǎn)生 16,631 個(gè)訓(xùn)練樣本、1,516 個(gè)驗(yàn)證樣本和 1,864 個(gè)測(cè)試樣本。

為了評(píng)估泛化能力，研究人員采用時(shí)間劃分策略，考慮到一些基線使用預(yù)訓(xùn)練的 ESM2 模型，存在數(shù)據(jù)泄漏的風(fēng)險(xiǎn)。時(shí)間劃分評(píng)估將特定日期之前的數(shù)據(jù)分配給訓(xùn)練集，而將該日期之后的數(shù)據(jù)分配給測(cè)試集。對(duì)于結(jié)構(gòu)的時(shí)間劃分評(píng)估，使用 CASP15 數(shù)據(jù)集，其中包含在訓(xùn)練期間未見過的新晶體結(jié)構(gòu)；對(duì)于序列的時(shí)間劃分評(píng)估，使用 NovelPro 數(shù)據(jù)集，該數(shù)據(jù)集包含 2023 年 11 月 23 日之前 30 天內(nèi)發(fā)布的 76 個(gè)蛋白質(zhì)序列，結(jié)構(gòu)由 AlphaFold 2 預(yù)測(cè)。
**RNA 設(shè)計(jì)任務(wù)中，**研究人員在 RDesign 收集的數(shù)據(jù)集上進(jìn)行 RNA 實(shí)驗(yàn)，該數(shù)據(jù)集包含 2,218 個(gè) RNA 三級(jí)結(jié)構(gòu)，這些結(jié)構(gòu)根據(jù)其結(jié)構(gòu)相似性分為訓(xùn)練集 (1,774 個(gè)結(jié)構(gòu))、測(cè)試集 (223 個(gè)結(jié)構(gòu)) 和驗(yàn)證集 (221 個(gè)結(jié)構(gòu))。由于數(shù)據(jù)樣本數(shù)量較少，研究人員報(bào)告了 3 次獨(dú)立運(yùn)行的中位數(shù)恢復(fù)率及其標(biāo)準(zhǔn)差。
材料設(shè)計(jì)任務(wù)中，研究人員在 CHILI-3K 數(shù)據(jù)集上評(píng)估 UniIF，該數(shù)據(jù)集由單金屬氧化物衍生的納米材料圖構(gòu)成。數(shù)據(jù)集包括 53 種金屬元素和一種非金屬元素 (氧)，共計(jì) 3,180 個(gè)圖，6,959,085 個(gè)節(jié)點(diǎn)和 49,624,440 條邊。
模型架構(gòu)：用于一般分子逆折疊的統(tǒng)一模型 UniIF
如下圖所示，研究人員提出了一個(gè)用于一般分子逆折疊的統(tǒng)一模型。

① 該模型將所有類型的分子 (All Molecules) 轉(zhuǎn)換為塊圖——對(duì)于大分子 (Macromolecules)，使用基于氨基酸和核苷酸的預(yù)定義框架；對(duì)于小分子 (Small molecules)，通過一層 GNN 學(xué)習(xí)每個(gè)塊的局部框架；

② 使用幾何特征提取器 (Geometric Featurizer) 初始化幾何節(jié)點(diǎn)特征 (Node feature) 和邊特征 (Edge features)；

③ 提出了塊圖注意力層 (Block Graph Attention)，基于此構(gòu)建塊圖神經(jīng)網(wǎng)絡(luò) (Block Graph Neural Network) 以學(xué)習(xí)表達(dá)豐富的塊表示；

④ 最后展示 UniIF 在多種任務(wù) (Tasks) 上均能取得具有競(jìng)爭(zhēng)力的結(jié)果，包括蛋白質(zhì)設(shè)計(jì)、RNA 設(shè)計(jì)和材料設(shè)計(jì)。

圖：UniIF 模型整體框架

**構(gòu)建塊圖：**該模型架構(gòu)的第一步是引入塊圖來表示所有類型的分子，關(guān)鍵在于將不規(guī)則的原子集合 (大小各異) 轉(zhuǎn)換為規(guī)則的塊表示 (固定大小)。研究人員引入基于框架的塊表示，以統(tǒng)一對(duì)所有分子的建模，一個(gè)塊包含等變框架和不變特征向量，局部框架包含軸矩陣和位移向量。對(duì)于大分子，軸矩陣是基于氨基酸和核苷酸預(yù)定義的；而對(duì)于小分子，因?yàn)樾》肿記]有先驗(yàn)的共同結(jié)構(gòu)模式，需要學(xué)習(xí)軸矩陣。給定一個(gè)包含 n 個(gè)塊的分子，研究人員使用 kNN 算法構(gòu)建塊圖。

圖：不同分子的結(jié)構(gòu)單元，基本構(gòu)建單元包括氨基酸、核苷酸和原子

**塊圖特征提取：**對(duì)于小分子，無法使用預(yù)定義的局部框架，因此研究人員需要為每個(gè)原子學(xué)習(xí)局部框架——即給定一個(gè)分子 ，其使用一層 GNN 來初始化原子表示，然后使用幾何特征提取器初始化幾何節(jié)點(diǎn)特征和邊特征。

塊圖注意力模塊：研究人員引入了幾何塊注意力網(wǎng)絡(luò)，包括幾何交互、交互注意力和虛擬長(zhǎng)期依賴模塊，以捕捉所有分子的三維交互。
研究結(jié)果：UniIF 在所有任務(wù)上都優(yōu)于最先進(jìn)的方法
研究人員通過多個(gè)逆折疊任務(wù)和消融研究展示了 UniIF 的有效性，包括：
* 蛋白質(zhì)設(shè)計(jì) (T1)：設(shè)計(jì)能夠折疊成目標(biāo)結(jié)構(gòu)的蛋白質(zhì)序列
* RNA 設(shè)計(jì) (T2)：設(shè)計(jì)能夠折疊成目標(biāo)結(jié)構(gòu)的 RNA 序列
* 材料設(shè)計(jì) (T3)：從已知材料結(jié)構(gòu)中發(fā)現(xiàn)穩(wěn)定的組成
① 蛋白質(zhì)設(shè)計(jì) (T1)

蛋白質(zhì)設(shè)計(jì)旨在設(shè)計(jì)能夠折疊成目標(biāo)結(jié)構(gòu)的蛋白質(zhì)序列，研究人員在不同設(shè)置 (有和沒有 ESM2) 和多個(gè)數(shù)據(jù)集 (CATH4.3、CASP、NovelPro) 下提供了結(jié)果。如下表所示：使用不包含 ESM2 的純逆折疊模型，UniIF 在所有數(shù)據(jù)集上實(shí)現(xiàn)了最佳性能，證明了其有效性。

*LMDesign 和KWDesign 包含ESM2；StructGNN、GraphTrans、GCA、GVP、AlphaDesign、ProteinMPNN 和 PiFold 不包含 ESM2

在 CATH4.3 上，由于基線模型較強(qiáng)，整體提升有限，但時(shí)間劃分評(píng)估突顯了 UniIF 在泛化能力上的優(yōu)勢(shì)，UniIF 以更少的可學(xué)習(xí)參數(shù)超越了強(qiáng)基線 PiFold。在時(shí)間劃分評(píng)估中，UniIF 以顯著的優(yōu)勢(shì)超過了所有基線，包括基于 ESM2 的方法。在包含新序列的 NovelPro 上，UniIF 的表現(xiàn)優(yōu)于使用 ESM2 進(jìn)行序列優(yōu)化的 LMDesign 和 KWDesign——這表明 UniIF 具有優(yōu)越的泛化能力，對(duì)于實(shí)際應(yīng)用至關(guān)重要。

蛋白質(zhì)設(shè)計(jì)結(jié)果 (最佳和次優(yōu)結(jié)果用加粗和下劃線標(biāo)注)
② RNA 設(shè)計(jì) (T2)

RNA 設(shè)計(jì)的目標(biāo)是設(shè)計(jì)能折疊成目標(biāo)結(jié)構(gòu)的 RNA 序列。如下表所示，UniIF 在所有情況下均取得最佳性能，這一提升是顯著的，因?yàn)橹皬?qiáng)大的基線模型如 PiFold 僅在蛋白質(zhì)設(shè)計(jì)上表現(xiàn)突出。據(jù)悉，UniIF 是第一個(gè)在蛋白質(zhì)和 RNA 設(shè)計(jì)任務(wù)中都實(shí)現(xiàn)最先進(jìn)性能的模型，證明了其多功能性和有效性。

RNA設(shè)計(jì)的恢復(fù)結(jié)果 (最佳和次優(yōu)結(jié)果用加粗和下劃線標(biāo)注)

③ 材料設(shè)計(jì) (T3)

從已知材料結(jié)構(gòu)中發(fā)現(xiàn)穩(wěn)定的原子組合對(duì)于新材料的發(fā)現(xiàn)至關(guān)重要，所以研究人員也評(píng)估了 UniIF 在這一新任務(wù)上的表現(xiàn)。如下表所示，UniIF 顯著超越所有基線模型。

表：CHILI-3K 結(jié)果

④ 案例研究

在下圖中，研究人員展示了設(shè)計(jì)的蛋白質(zhì)和 RNA 序列。此外，其使用 AlphaFold 3 將設(shè)計(jì)的序列重新折疊成結(jié)構(gòu)——真實(shí)結(jié)構(gòu)（灰色）、PiFold 結(jié)構(gòu)（綠色）和 UniIF 結(jié)構(gòu)（粉色）進(jìn)行了對(duì)齊和比較。研究人員觀察到，UniIF 在恢復(fù)率和均方根偏差 (RMSD) 方面都取得了改善，證明了其在逆折疊任務(wù)中的有效性。

圖：設(shè)計(jì)示例，將真實(shí)結(jié)構(gòu) (灰色)、PiFold 結(jié)構(gòu) (綠色) 和 UniIF 結(jié)構(gòu) (粉色) 對(duì)齊
UniIF 模型對(duì) AlphaFold 3 形成進(jìn)一步補(bǔ)充
通用分子學(xué)習(xí)在近年來受到越來越多的關(guān)注，RoseTTAFold All-Atom (RFAA) 和 AlphaFold 3 是兩個(gè)在該方向取得顯著成功的代表性模型。

2024 年 3 月 7 日，David Baker 在 Science 發(fā)布了題為「Generalized biomolecular modeling and design with RoseTTAFold All-Atom」的研究論文。該團(tuán)隊(duì)開發(fā)了 RoseTTAFold All-Atom (RFAA)，它可以將氨基酸和 DNA 堿基基于殘基的表示與所有其他基團(tuán)的原子表示相結(jié)合，從而對(duì)包含蛋白質(zhì)、核酸、小分子、金屬和給定序列和化學(xué)結(jié)構(gòu)的共價(jià)修飾組件進(jìn)行建模。

2024 年 5 月 9 日，Demis Hassabis、John Jumpe 等人在 Nature 發(fā)表了題為「Accurate structure prediction of biomolecular interactions with AlphaFold 3」的研究論文。該研究推出了 AlphaFold 3，這一最新模型能預(yù)測(cè)含有蛋白質(zhì)數(shù)據(jù)庫 (Protein Data Bank) 內(nèi)幾乎所有分子類型的復(fù)合物的結(jié)構(gòu)，包括配體 (小分子)、蛋白質(zhì)、核酸 （DNA 和 RNA) 如何聚集在一起并相互作用，以及預(yù)測(cè)翻譯后修飾和離子對(duì)這些分子系統(tǒng)的結(jié)構(gòu)影響，從而幫助科研人員在原子水平上精確地觀察生物分子系統(tǒng)的結(jié)構(gòu)。

細(xì)究這兩種模型，RFAA 使用原子-鍵圖表示小分子，而使用框架圖表示大分子；AlphaFold 3 則采用雙層表示，即原子表示和標(biāo)記表示，適用于所有分子。標(biāo)記概念就相當(dāng)于前文所述的塊概念，表示一組原子，如氨基酸或核苷酸。

GET 和 EPT 是最近提出的兩個(gè)模型，采用塊表示法同時(shí)適用于小分子和大分子，并引入了新的等變變換器骨架。與指定小分子原子-鍵圖的 RFAA 不同，本文介紹的 UniIF 模型為所有分子類型采用統(tǒng)一塊圖，不需要原子-鍵圖，而且該模型還為每個(gè)塊引入了向量基，這一點(diǎn)與 AlphaFold 3 、GET 和 EPT 都不同。

由于在一定程度上解決了構(gòu)建通用分子模型的挑戰(zhàn)，**UniIF 模型可以視作在 RoseTTAFold All-Atom 和 AlphaFold 3 等「前輩們 」分子結(jié)構(gòu)預(yù)測(cè)方向取得進(jìn)展的進(jìn)一步補(bǔ)充。**未來，不斷迭代的生物大模型將幫助研究人員重新認(rèn)識(shí)生物世界、重新思考藥物發(fā)現(xiàn)，從而造福于全人類。