機器之心發布
機器之心編輯部近日,百圖生科與清華大學聯合提出了一種名為 xTrimo Protein General Language Model (xTrimoPGLM) 的模型,參數量高達千億(100B)。相關成果於 2023 年 7 月 7 日在 biorxiv 上發布。自然語言處理(NLP)領域中的預訓練模型相關進展已經被成功地用於學習蛋白質序列中隱藏的生物信息。現在大多數的蛋白質預訓練模型都受限於自動編碼或自回歸預訓練目標,這使得它們難以同時處理蛋白質理解(例如,蛋白質結構預測)和生成任務(例如,藥物設計)。這篇論文提出統一的蛋白質語言模型,xTrimoPGLM,通過一個創新的預訓練框架來同時處理這兩種類型的任務。xTrimoPGLM 主要技術貢獻是探索了這兩種類型目標之間的兼容性以及共同優化的可能性,並基於此訓練了一個前所未有的 1000 億參數規模的蛋白質語言模型,並消耗了 1 萬億 Tokens,模型 FLOPs 達到 6.2e+23,達到和 175B 參數的 GPT-3 一個量級。在理解任務上,xTrimoPGLM 在多種蛋白質理解任務(15 項任務中的 13 項任務)中顯著優於其他先進基線。在生成任務上,xTrimoPGLM 能夠生成與自然蛋白質結構類似的新蛋白質序列。此外,文章基於相同的框架額外訓練了一個 12 億參數的抗體模型(xTrimoPGLM-Ab),其在預測抗體自然性和結構方面取得了市面上最好的效果,並且顯示出比 AlphaFold2 更快的推理速度(數十倍到數千倍)。綜合來看,這些結果充分展示了 xTrimoPGLM 在理解和生成蛋白質序列方面的強大能力和廣闊的應用前景。蛋白質理解和生成的統一不同類型的蛋白質相關任務需要蛋白質語言模型(Protein Language Model,PLM)提供各異的輸出。具體的,蛋白質理解任務,如二級結構預測等,需要 PLM 提供精確的氨基酸和序列級別的表示;而蛋白質設計任務,如抗體或酶的設計,依賴於 PLM 的生成能力。然而,當前的 PLM 因為其單一的預訓練框架的限制,大多只能處理一種類型的任務。事實上,蛋白質的理解和生成都反映了蛋白質數據的分布信息,Meta 之前使用 ESM(基於 Masked Language Model 的蛋白質大模型)做生成的工作也證實了這一點,指出蛋白質理解預訓練模型可以通過一些採樣策略進行蛋白質設計。這進一步支撐了這兩種看似不同的任務的統一性,如果能夠使用同一個訓練框架去處理這兩種任務,將會進一步增強模型對蛋白質數據的擬合能力。雖然 NLP 領域生成式模型(例如 UL2R) 已經成為主流範式,通過把數據標籤映射到整個文本空間,結合指令微調來生成各式各樣的任務的答案,但 PLM 還無法實現這一點。實際上,蛋白質的應用仍然依賴於將表示與下游任務特定標籤(如結構預測的 3D 坐標)之間的橋接,這在很大程度上依賴於 BERT 樣式的訓練來處理蛋白質理解任務。因此,需要同時進行這兩種訓練目標。本文介紹的 xTrimo 蛋白質通用語言模型(xTrimoPGLM)預訓練框架,巧妙地統一了兩種類型的預訓練任務,使模型能同時處理多種與蛋白質相關的任務。研究團隊通過全面的實驗評估了 xTrimoPGLM 框架的有效性。在蛋白質理解任務的情況下,xTrimoPGLM-100B 在多種評估中表現出色,涵蓋了蛋白質結構、功能、交互和可開發性等領域的 15 項任務。"Scaling Law" 是衡量大型語言模型的重要原則,模型的性能應隨着模型參數大小、數據量、計算量按比例指數增加而線性增長。實際上,研究團隊在下游任務上的實驗結果驗證了這一定律,證明了大型模型在處理複雜任務時的必要性。如下圖所示,性能改善與預訓練計算量增加之間的關係。以 Meta 的 ESM-2 為參考,隨着蛋白質語言模型(PLM)的計算量呈指數增長,蛋白質的下游性能仍然會線性增長(每個大類有 3-4 個任務,數值表示這些任務的平均值)。 在蛋白質生成任務中,xTrimoPGLM-100B 展示了生成不同長度和序列的新蛋白質序列的能力,這是通過調整生成超參數實現的。值得注意的是,當與自然蛋白質進行比較時,xTrimoPGLM-100B 展示了生成結構相似但序列相異的新蛋白質序列的能力。這再次驗證了大型模型對於複雜任務的重要性,進一步證實了遵循 "Scaling Law" 進行模型設計的決定是正確的。同時,研究團隊還開發了一種當前更具實用意義的具有 12 億參數的抗體 PLM,即 xTrimoPGLM-Ab-1B。這種模型在 OAS 抗體數據庫上進行 fine-tuning,處理了超過 1 萬億個 token。它在抗體的自然性和結構預測任務上達到了目前最優秀的性能。由於不依賴於外部庫的檢索和多序列對齊(Multiple Sequence Alignment),所以結構預測在速度上比 Alphafold2 模型提升成百上千倍,這對於基於抗體藥物發現 AI 製藥公司至關重要。混合訓練原始的 GLM 模型利用兩種不同的預訓練目標來提升其生成能力:1)跨度生成(Span Generation,簡稱 sMask),用於恢複句子中的短空白;2)長文本生成(Long-text Generation,簡稱 gMASK),用於在提供前綴上下文的基礎上生成隨機長度的序列。為了進一步提升 xTrimoPGLM 的理解能力,團隊在 prefix 區域引入了被用作理解目標的 Masked Language Model(MLM,即 [MASK])。這樣的設計確保了 xTrimoPGLM 能生成精確的殘基級和序列級表示。當使用 [MASK] 標識符時,xTrimoPGLM 的功能類似於 BERT。相反,當使用 [sMASK] 或 [gMASK] 時,xTrimoPGLM 的行為類似於 PrefixLM 或 GPT。總的來說,xTrimoPGLM-100B 的預訓練階段可以分為兩個階段。首先,利用 MLM 進行預訓練以增強其表示能力,主要目標是快速減少損失水平。第二階段,使用結合 MLM 和 GLM 損失的統一目標進行訓練,以提升理解和生成能力。 NLP 領域大量探索了統一的預訓練模式,但大多還是採樣了同樣的訓練模式(自回歸或自編碼)。為了滿足統一的蛋白質預訓練模型的需求,需要將 BERT 樣式的目標引入到預訓練語言模型中,以增強模型的表示能力,同時也需要引入 GPT 樣式的目標,以確保模型的生成能力。在最開始研究團隊使用 Probing 策略探索 Contact Map Prediction 的任務時,發現僅僅依靠基於下一個詞預測的生成式語言模型,效果會有大幅度的下降。兼容性實驗:在實證分析中,研究團隊在 xTrimoPGLM-150m 模型上探究了同時優化兩個不同目標的可行性。結果如下圖 (a)(b) 所示,即使兩種預訓練目標看似衝突,MLM 損失和 GLM 損失也可以同時優化,反之亦然,即統一的訓練可以很快的適配到 MLM 或者 GLM 上,並不會影響兩者的收斂情況。研究團隊還探究了 MLM 與 GLM 兩種目標是否能相互加速收斂,結果如圖 (c)(d)。
- MLM-adapted GLM:接着 MLM 預訓練後的模型,繼續訓練 GLM 目標函數;
- GLM-adapted MLM:接着 GLM 預訓練後的模型,繼續訓練 MLM 目標函數;
- CDR3 短序列掩碼(CSM):對 CDR3 部分區域進行掩碼重設計。
- CDR3 全序列掩碼(CWM):對 CDR3 全部區域進行掩碼重設計。
- CDR3 隨機突變(CRM):對 CDR3 區域內特定位點的隨機 3-6 個位置進行突變。
- CDR3 隨機檢索(CRR):使用來自 SARS-CoV-2 野生型庫中其他抗體的序列隨機替換 CDR3 區域。
© THE END
轉載請聯繫本公眾號獲得授權