目前,蛋白質基座模型領域的相關研究仍停留在無法充分適配蛋白質序列生物學特性的「BERT」時代。此前,AlphaFold 和 ESM 等 AI 模型顯著推動了包括結構預測、逆向折疊、功能特性預測、突變效應評估以及蛋白質設計等多個領域的發展。但這些模型仍缺乏類似于前沿大語言模型(LLMs)的可擴展、系統化方法論,其能力無法實現隨數據量、模型規模和計算資源的增加而持續提升。
此類模型的通用性缺陷,帶來了蛋白質設計領域難以解決的挑戰:模型無法捕捉蛋白質的構象異質性,對蛋白質設計的預測也無法超越訓練數據的范疇;而對 NLP 方法論遷移的過度依賴,則導致了針對蛋白質特性的原創性架構設計的缺失。
在此背景下,清華大學智能產業研究院(AIR)周浩課題組聯合上海人工智能實驗室,基于貝葉斯流網絡提出經系統化訓練的蛋白質基座模型 AMix-1,為蛋白質設計提出了一種可擴展的通用路徑。該模型首次采用了「Pretraining Scaling Law」、「Emergent Ability」、「In-Context Learning」以及「Test-time Scaling」的系統化方法論,并在這一基礎上設計了一種基于多序列比對(MSA)的上下文學習策略,實現了蛋白質設計通用框架的一致性,同時確保了模型的可擴展性。
相關研究成果以「AMix-1: A Pathway to Test-Time Scalable Protein Foundation Model」為題,發布于 arXiv 平臺。
研究亮點:
* 為基于貝葉斯流網絡的蛋白質生成模型建立了可預測的 Scaling Law;
* AMix-1 模型僅通過序列級訓練目標,自發涌現出對蛋白質結構的「感知理解」能力,無需顯式結構監督;
* 基于多序列比對(MSA)的上下文學習框架解決了功能優化中的對齊難題,升級了模型在進化語境中的推理和設計能力,使 AMix-1 能夠生成結構和功能保守的新蛋白質;
* 提出驗證成本導向的測試時擴展算法,在驗證預算增加時開啟基于演化的設計新方式。
UniRef50 數據集:預處理與迭代聚類
研究人員在模型預訓練階段使用了預處理的 UniRef50 數據集。該數據集由 EvoDiff 提供,源自 UniProtKB,通過迭代聚類從 UniParc 序列中篩選得出(UniProtKB+UniParc → UniRef100 → UniRef90 → UniRef50),包含 41,546,293 條訓練序列和 82,929 條驗證序列。其中,對于長度超過 1,024 個殘基的序列,研究人員使用隨機裁剪的策略,將其長度裁減至 1,024 個殘基以降低計算成本并生成多樣化子序列。該迭代流程確保了 UniRef50 代表性序列高質量、非冗余且多樣化的特性,為蛋白質語言模型提供了覆蓋廣泛的蛋白質序列空間。
系統性技術方案
AMix-1 為蛋白質基座模型實現 Test-time Scaling 提供了一整套系統性的技術方案:
* Pretraining Scaling Law:明確了參數、樣本數和計算量如何權衡,才能最大化模型的能力;
* Emergent Ability:顯示隨著訓練的推進,模型會涌現出對蛋白結構的「感知理解」;
* In-Context Learning:解決了功能優化中的對齊難題,讓模型學會在進化語境中推理與設計;
* Test-time Scaling:AMix-1 在驗證預算增加時開啟基于演化的設計新方式。
從訓練、推理到設計,AMix-1 論證了其作為蛋白質基座模型的通用性和可擴展性,為走向實際落地做好了鋪墊。
模型能力對比
Pretraining Scaling Law:可預測的蛋白質模型能力
該研究為實現 AMix-1 的可預測 Scaling Law,在實驗中設計了從 800 萬到 17 億參數的多尺度模型組合,并利用訓練浮點操作數(FLOPs)作為統一衡量指標,精確擬合、預測了模型交叉熵損失與計算量的冪律關系。
從結果表征來看,模型損失與計算量之間的冪律曲線吻合度極高,證實了基于貝葉斯流網絡的模型訓練過程具有高度可預測性。
不同噪聲尺度下 AMix-1 的 Scaling Law
Emergent Ability:實現模型高級能力躍升
在蛋白質序列學習中,結構信息的涌現研究通常基于「序列-結構-功能」的范式。為驗證蛋白質建模中優化動力學與功能結果的關聯,研究團隊基于可預測的 Scaling Law 采用以損失函數為核心的視角分析涌現行為,在將預測性交叉熵損失作為錨點的同時,通過經驗方法將訓練損失映射至蛋白質的生成性能。該研究對模型 Emergent Ability 的評估聚焦于 3 個層面:
* 從序列一致性觀照模型從受損序列分布中恢復序列水平的能力;
* 從可折疊性看模型從序列理解到結構可行性的轉變;
* 從結構一致性判斷模型保持結構特征的能力。
AMix-1 在訓練中的相關數據完整展現了蛋白質基座模型「序列一致性、可折疊性、結構一致性」能力的涌現過程。數據顯示,模型在訓練過程中的所有能力指標與交叉熵損失呈高度相關性,驗證了通過 Scaling Law 和交叉熵損失預測模型能力的可能性。同時,在僅以序列級自監督目標訓練、完全不引入結構信息的情況下,模型依然在交叉熵損失下降至閾值后展現出 Emergent Ability,表現出 pLDDT 與 TM-score 的非線性躍遷。
AMix-1 模型在不同指標下的涌現現象
In-Context Learning:蛋白質設計的通用范式
通過計算機模擬案例,研究人員驗證了 AMix-1 的 In-Context Learning 機制。模擬案例實驗顯示,AMix-1 能夠從輸入樣本中準確提取并泛化結構或功能約束,而無需依賴顯式標簽或結構監督。
LLM 與 AMix-1 在 In-Context Learning 層面的比較
相較于需要根據任務類型進行定制流程、缺乏統一的蛋白設計框架的傳統蛋白質設計,AMix-1 引入了大語言模型中的上下文學習(ICL)機制,以完成結構和功能導向的蛋白設計。實驗顯示,AMix-1 在結構任務上,能夠以常規同源蛋白甚至在近乎無同源的蛋白為提示生成在預測結構上高度一致的新蛋白,在功能任務中則能在輸入蛋白的酶學功能和化學反應引導的酶設計上,生成出功能高度一致的蛋白酶。
在這一通用機制下,模型能夠在給定的一組蛋白質中,自動推理其中的共性信息和規律,并運用此規律來指導生成新的符合共性規律的蛋白。該機制將一組蛋白質 MSA 壓縮成一個位置級的概率分布(Profile)輸入模型,模型在快速辨析所輸入蛋白質的結構及功能規律后,即可生成符合意圖的新蛋白。
驗證 AMix-1 In-Context Learning 的案例
Test-time Scaling:可擴展的通用智能
基于 Test-time Scaling 方法,研究人員采用提議者-驗證者框架(Proposer-Verifier Framework)構建了 EvoAMix-1,通過持續增加驗證預算提升了 AMix-1 的模型性能。研究團隊在放大模型的設計效能的同時,實現了模型的可擴展性。同時,為確保兼容性,研究團隊減去了對驗證者的性質預設。
EvoAMix-1 基于概率模型固有的隨機性促進探索,通過整合特定任務的計算機模擬獎勵函數或實驗檢測反饋,在進化約束下迭代生成并篩選候選蛋白質序列,無需模型微調即可實現高效的定向蛋白質進化,在蛋白質設計中實現了穩健且測試時可擴展的性能。在全部 6 項設計任務中,EvoAMix-1 在 In-Context Learning 及各類強基線方法上始終優于 AMix-1 。
AMix-1 在推理階段進行進化擴展算法的工作流
和通過重要性采樣生成新蛋白質變體的傳統方法相比,EvoAMix-1 并不更新模型參數,而是通過上下文范例構建提議分布。在每一輪中,AMix-1 都會以一組多序列比對(MSA)或其譜作為提示,將其視為蛋白質基礎模型的輸入條件,隨后基礎模型對鄰近序列進行采樣,從而有效定義新的條件提議分布。
研究團隊在多個代表性蛋白質定向進化任務上系統驗證了 EvoAMix-1 的通用性和可擴展性,包括酶的最適 pH 和溫度進化、功能保持與增強、孤兒蛋白設計,以及通用的結構引導優化等。實驗結果顯示,EvoAMix-1 Test-time Scaling 的擴展能力強勁,在跨任務、跨目標的情況下仍表現出強大的通用性。
EvoAMix-1 在六個定向進化基準中的 Test-time Scaling 性能
濕實驗驗證:AMix-1 輔助研發蛋白 AmeR 變體活性提升 50 倍
該研究在實際濕實驗中檢驗了「上下文提示設計」策略,進一步驗證了 AMix-1 在高效設計高活性 AmeR 變體方面的優勢。研究人員選取目標蛋白 AmeR,使用 AMix 模型基于 AmeR 家族的概率分布生成 40 個變體,并通過熒光報告基因實驗評估每個變體的抑制能力。其中,每個變體僅包含 ≤10 處氨基酸突變,倍數抑制(Fold Repression)的數值越高功能越強。此外,該研究還提出了進化測試時的縮放算法,以增強 AMix-1 在蛋白質定向進化中的適用性,并通過多種計算機模擬靶區指標驗證其性能。
最終結果顯示,由 AMix-1 產生的最優變體活性提升高達 50 倍,其效能與當前 SOTA 模型相比提升了約 77% 。此外,AMix-1 不依賴反復篩選或手動設計,完全由模型自動生成,實現了「從模型到實驗」的完整閉環,同時達成了 AI 在功能蛋白設計上的首次量級突破。
AMix-1 模型在濕實驗中的生成迭代和相應的變異過程
全局拓撲,感知開啟蛋白質設計的全新維度
目前,AI 與蛋白質設計的結合研究正方興未艾。除了 AMix-1,清華大學生命學院龔海鵬課題組聯合北京生命科學研究所徐純福研究團隊提出的幾何感知擴散模型 TopoDiff,同樣在蛋白質設計領域取得了重大突破。
RFDiffusion 等傳統擴散模型在生成免疫球蛋白等特定折疊類型時,不僅存在覆蓋偏差,而且缺乏對蛋白質全局拓撲的量化評估指標。該研究基于 CATH 、 SCOPe 等結構數據庫,提出了無監督系統 TopoDiff 框架,通過學習和利用全局幾何感知的潛在表征,實現了基于擴散模型的無條件與可控蛋白質生成。該研究提出「覆蓋度」(Coverage)的新評估指標,通過編碼器-擴散模型雙階段框架,將蛋白質結構解耦為全局幾何藍圖與局部原子坐標生成,突破了蛋白質折疊覆蓋的研究難題。
此外,NVIDIA 聯合加拿大魁北克人工智能研究所 Mila,還基于 AlphaFold 架構升級的全原子生成模型則攻克了長鏈預測難題。傳統方法不僅難以生成超長鏈(>500 殘基)的全原子結構,而且無法探索膜蛋白特異性口袋等非經典折疊構象。研究團隊引入的概率決策機制,以量子場論中的路徑積分采樣替代了確定性折疊軌跡,從而使膜蛋白設計成功率提升至 68% 。
幾何感知解構蛋白質折疊,攻克 500+ 殘基長鏈設計,自然語言指令驅動蛋白設計,靶向「不可成藥」IDPs 蛋白……AI 在拓展蛋白質設計能力邊界的同時,為蛋白質設計領域的研究路徑提供了新的范式。未來,AI 驅動的蛋白質設計有望為開發創新療法、酶制劑和生物材料開拓更廣闊的空間。
來源: HyperAI超神經