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

分子逆折疊在藥物和材料設(shè)計中起到關(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ī)模：小分子允許全局注意力機制來學(xué)習長期依賴關(guān)系，但這往往在大分子上行不通。

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

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

研究亮點：

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

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

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

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

蛋白質(zhì)設(shè)計任務(wù)中**，研究人員在 CATH4.3 數(shù)據(jù)集上評估 UniIF。該數(shù)據(jù)集按 CATH 拓撲分類代碼分割，產(chǎn)生 16,631 個訓(xùn)練樣本、1,516 個驗證樣本和 1,864 個測試樣本。

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

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

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

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

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

圖：UniIF 模型整體框架

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

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

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

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

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

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

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

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

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

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

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

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

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

④ 案例研究

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

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

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

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

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

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

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