中科院計算所團隊提出CarbonNovo，基於AI進行蛋白質結構和序列的端到端從頭設計

編輯 | ScienceAI

作者 | 計算所張海倉團隊

近期，中國科學院計算所張海倉帶領的研究團隊提出了 CarbonNovo，以端到端的方式聯合設計蛋白質主鏈結構和序列。

該研究以「CarbonNovo: Joint Design of Protein Structure and Sequence Using a Unified Energy-based Model」為題發表在機器學習會議 ICML 2024 上。

背景介紹

蛋白質是生物執行功能的重要大分子。蛋白質從頭設計旨在創造全新的蛋白質，在藥物開發和酶工程中有著廣泛的應用。

近幾年，基於 AI 的蛋白質從頭設計快速發展，已被成功應用於抗體設計、小蛋白藥物設計等領域，和傳統設計方法相比，其顯著提高了設計成功率和效率。

AI 蛋白質設計得益於近幾年的兩大技術突破：

一是蛋白質結構預測領域的 AlphaFold2 模型，它為蛋白質計算領域包括蛋白質設計領域，提供了基礎的神經網路模型架構、蛋白質序列表示和結構表示方法、以及先進的訓練策略（蒸餾訓練、端到端訓練）等技術；

二是 AIGC 在文字、影像、影片生成領域的快速發展，為蛋白質設計提供了成熟的生成模型，例如 DDPM, SDE, Flow Matching，Bayesian Flow Network 等。代表性的蛋白質設計模型，例如 RFDiffusion 和 Chroma 等，主要思路都是將這兩大技術融合，將蛋白質的序列和結構表示網路嵌入到基於 AI 的生成模型框架中。

圖 1：蛋白質從頭設計的「兩階段」框架。(來源，作者)

蛋白質從頭設計主要包含兩個步驟，蛋白質主鏈結構設計和序列設計（圖1）。相應地，當前的主流模型通常採用「兩階段」框架進行：在訓練過程中，結構設計模組和序列設計模組分別訓練；在推斷過程中，首先生成主鏈結構，然後為主鏈結構生成最優序列。代表性地，領域內使用 RFDiffusion 和 ProteinMPNN 這兩個軟體，依次生成主鏈結構和序列。

蛋白質從頭設計的「兩階段」框架存在固有的侷限性：

(1) 序列設計模組面臨過擬合風險。序列設計模組是在準確的結晶結構上訓練的，而在推斷階段，結構模組生成的結構是包含噪音的，其準確性與晶體結構不一致。

(2) 序列設計模組缺乏與結構設計模組的互動，設計序列無法為結構生成模組提供反饋以進一步最佳化生成的結構。

CarbonNovo進行端到端的結構和序列聯合設計

針對蛋白質設計「兩階段」框架的侷限性，中科院計算所張海倉帶領的研究團隊提出 CarbonNovo，以端到端的方式聯合設計蛋白質主鏈結構和序列。論文近期已於近期發表在機器學習會議 ICML 2024 上。

圖 2：CarbonNovo 端到端生成蛋白質結構和序列。（來源：論文）

CarbonNovo 的主要貢獻總結如下：

1）設計了基於能量的生成模型，以端到端的方式設計主鏈結構和序列（圖 2）。之前的結構-序列聯合設計模型主要針對抗體等特殊的蛋白質家族，CarbonNovo 是第一個針對所有蛋白質家族的結構-序列聯合設計模型。

2）基於 network recycling 技術，CarbonNovo 首次將蛋白質語言模型引入蛋白質結構設計任務，以利用海量天然蛋白質序列資料包含的先驗資訊。

3）CarbonNovo 採用多種技術，以提高結構-序列聯合生成模型的訓練和推斷效率，例如多階段訓練策略和用於序列取樣的離散 M-H Langevin 演算法。

蛋白質結構-序列的聯合能量模型

在經典物理模型下，天然蛋白質構象具有比較低的自由能，這也是蛋白質結構預測和設計的一般假設。基於此，CarbonNovo 建立了蛋白質結構和序列的聯合能量模型:

1) 採用SE(3)空間的擴散模型刻畫主鏈結構能量。最近的一些研究，無論是在訓練目標還是推理策略方面，都已將擴散模型統一在能量模型（Energy-based model）下。

CarbonNovo也是基於此，將AI模型下的能量和經典物理模型下的能量聯絡在一起。需要指出的是，主流擴散模型（例如 FrameDiff 和 Genie）只是用IPA (Invariant Point Attention network)作為生成網路。而 CarbonNovo 引入了 AlphaFold2 模型裡的三角更新網路，也是該工作的創新之一。

2) 採用 amortized Potts 模型刻畫給定主鏈結構下的序列能量。Potts 模型是經典的能量模型，在蛋白質結構預測領域被用來學習同源序列中的共進化訊號。CarbonNovo 利用學習到的結構表示來引數化 Potts 模型，建立給定結構條件的序列能量模型。

圖 2 展示了 CarbonNovo 的具體生成過程:

1) 基於擴散模型設計出當前時間步的主鏈結構。

2) 給定當前的主鏈結構，基於 Potts 模型設計出可能的序列。

3) 透過 network recycling 機制，將中間序列的語言模型表示回傳到結構模組來幫助生成更一致的序列-結構，實現了結構模組、序列模組和預訓練語言模型的統一。

CarbonNovo 生成蛋白質結構-序列的效能評測

圖 3：CarbonNovo 和「兩階段」設計模型的效能比較。（來源：論文）

論文采用多種指標充分評價了 CarbonNovo 在蛋白質從頭設計方面的效能（圖 3），例如，可摺疊性、多樣性、新穎性是領域內常用的評價指標。此外，論文還採用了 Rosetta 能量和語言模型下的似然機率（Sequence plausibility）作為評價指標。

CarbonNovo 和當前主流的「兩階段」設計模型做了比較，例如 RFdiffusion, Chroma, Genie, FrameDiff 和 FrameFlow。在最關鍵的可摺疊性指標上 CarbonNovo 顯著超過了所有基線方法，在其他指標上也顯著超過基線方法或者和基線方法相當。

為了展示 CarbonNovo 在聯合設計序列與結構方面的優勢，作者還對比了使用 ProteinMPNN 生成序列的結果（圖 3 a-c）。可以觀察到，聯合設計模型可以設計出更加匹配的蛋白質主鏈結構和序列。

圖 4：在不同蛋白質長度下的效能比較。（來源：論文）

作者進一步評估了 CarbonNovo 在不同長度的蛋白質設計上的效能（圖 4）。在設計比較短的蛋白質時（例如長度 100），各個模型的表現相當。而隨著蛋白質長度增長，CarbonNovo 的設計效能顯著優於「兩階段」設計模型。

消融實驗

圖 5：消融實驗結果。（來源：論文）

作者訓練了多個消融模型，以評估關鍵元件對 CarbonNovo 效能的相對貢獻（圖 5）。語言模型、序列設計模組和輔助訓練損失對 CarbonNovo 的效能都有貢獻。其中，語言模型的引入表現出最顯著的貢獻。此外，使用基於能量的序列設計模組相比自迴歸模型，也能顯著提升序列設計的效能。

Case study：蛋白質結構「插值」

圖 5：從 all beta-sheets 結構到 all alpha-helices 結構的漸變。（來源：論文）

在影像生成領域，人臉圖片插值/漸變是生成模型比較經典的應用。作者也嘗試利用 CarbonNovo 進行蛋白質結構的插值。

圖 5 展示了一個代表例子，隨著在隱空間逐漸增加 all alpha-helices 結構向量的權重，生成的 all beta-sheets 結構會逐漸過渡到 all alpha-helices 結構。

這是領域內第一個關於蛋白質結構的插值實驗，也體現了 CarbonNovo 學習到的蛋白質隱空間比較緊緻。

結語

最後作者指出，雖然 CarbonNovo 主要側重於蛋白質單體設計，但它也可以很容易被擴充套件到蛋白質複合物設計和條件設計，如多肽設計、抗體設計等。

作者團隊目前在跟生物實驗團隊合作，透過溼實驗的方式驗證 CarbonNovo 設計的蛋白質。

作者所在 CarbonMatrix 團隊，長期致力於 AI 蛋白質設計和 AI 藥物設計，正在建立生物大分子結構設計和預測的統一生成模型。

其研究成果多次發表在 ICML、NeurIPS 等頂級機器學習會議和 Nature Machine Intelligence、Nature Communications 等頂級學術期刊上，目前也在和生物實驗室合作，積極推動 AI 模型在藥物設計領域的產業化落地。

論文連結：https://openreview.net/pdf?id=FSxTEvuFa7

程式碼連結：https://github.com/zhanghaicang/carbonmatrix_public

相關文章