將蛋白質語言模型擴展到千億參數，深度解讀百圖生科、清華xTrimoPGLM模型
分類：綜藝
瀏覽數：4114
2023-07-18
機器之心發布
機器之心編輯部 近日，百圖生科與清華大學聯合提出了一種名為 xTrimo P rotein G eneral L anguage M odel (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 目標函數； 總的說來，基於預訓練後的模型中繼續訓練的模型，與從頭開始訓練的模型相比，其收斂速度明顯加快。這些實驗觀察到：在蛋白質數據分布並不依賴於特定的訓練模式，從而縮小了自編碼 PLMs（如 ESM）和自回歸 PLMs（如 ProGen2）之間的差距，為 100B 訓練流程提供了支持。
訓練穩定性 訓練穩定性是成功訓練 100B 規模大型語言模型的決定性因素 。xTrimoPGLM 從 GLM-130B 的實現中借鑒了一些想法並解決了許多不穩定訓練的問題。然而，xTrimoPGLM-100B 在從訓練的第一階段過渡到第二階段時仍然會遇到災難性的訓練崩潰（小規模的模型（10B 規模）中並未觀察到），即使一開始只將 1% 的 GLM 損失加入預訓練也可能觸發這些崩潰。 下圖可以看到，如果直接給 GLM 分配一個比例 ，在訓練初期，grad norm 都會出現 spike（橙，藍，綠線）。
為了緩解這個問題，研究團隊提出了一種平滑過渡策略。主要分為兩個階段。在第一階段，主要目標是逐步提高 GLM 損失的比例，達到預期的數量。具體來說，給定一個期望的 GLM 損失比例 R，按照線性增長，以 K 步從 0 增加到 R。在這個階段應該將學習率保持在極低的水平。完成過渡後，學習率可以按照預定義的腳本在幾百個 steps 內逐漸回升至原來的水平 （紫線）。實際上，最後的 xTrimoPGLM-100B 訓練運行只在過渡階段經歷了損失分歧情況，但是由於硬體故障多次失敗，導致經常性的更換節點和重啟。 訓練數據 為了能夠儘可能的映射整個蛋白世界， xTrimoPGLM-100B 的預訓練模型的訓練數據集整合自兩個廣泛的數據源：Uniref90 和 ColAbFoldDB。 結合這兩個數據源，預訓練模型數據集充分利用了這兩個數據源的優勢，既有廣泛的生物分類覆蓋，又有多樣的環境生態位蛋白質序列，全面而詳盡地映射了生物世界中的蛋白質資源。
超參數配置 xTrimoPGLM-100B 模型的訓練過程複雜且耗費資源，團隊耗費了 160 天的時間，開發團隊使用了 96 台 DGX-A100 GPU 伺服器（每台伺服器擁有 8×40G 的 GPU）以混合精度（FP16）進行訓練，消耗了 1 萬億的 tokens，由於大多數現有的大型語言模型在訓練上存在嚴重不足，所以開發團隊現在仍然在繼續訓練 xTrimoPGLM-100B 模型，以處理儘可能多的 tokens。 模型採用了 3D 並行策略，基於 DeepSpeed 進行了 4 路張量並行、8 路流水線並行，並且採用了 Zero Stage 1 進行訓練。模型具有 72 個 Transformer 層，80 個 head，以及 10,240 維和 31,744 個 FNN 的維度。使用了 DeepNorm 初始化 Post-LN，並且採用了 Embedding Layer Gradient Shrink (EGS) 穩定訓練，以及使用了 2D ROPE 位置編碼技術。為了提升訓練效率，每個樣本由多個蛋白拼接再一起，使用 < eos > 區分開來，包含固定的 2,048 的序列長度。最終單卡的 TFLOPs 在 120-135 之間，68 examples/sec, 如果是 80G 的 A100，經過減少重算可到 92 examples/sec。下表展示了大部分超參數的配置。
團隊也對比了當前比較熱門的預訓練模型的 FLOPs，目前生物領域不同模型的結構也存在比較大的區別， 所以在 FLOPs 的計算上，團隊考慮儘可能詳細統計，包括 query、 key、value 的轉換，Attention Matrix 的計算，注意力之後線性變換，MLP 中的變換，以及最後到 vocab 維度的映射，以及語言模型頭（如果存在的話）中的線性轉換等，可以看到，xTrimoPGLM-100B 高出其他模型一到兩個量級。
十五種蛋白質下游任務評估 為了全面評估 xTrimoPGLM-100B 模型，研究團隊對 4 個領域內，15 個下游的蛋白質相關任務進行了基準測試。測試結果表明，xTrimoPGLM-100B 在蛋白質結構、蛋白質可開發能力、蛋白質相互作用和蛋白質功能等四個主要類別的任務中，都表現出了顯著的優勢。在這些任務中，xTrimoPGLM-100B 模型與微調技術的組合取得了優異的成績，大部分超越了當前最先進的方法，從而推動這個領域的進步。 需要強調的是，下表對比主要從任務的角度進行，而不是一個完全公平的對比，因為 xTrimoPGLM-100B 在取得這些結果時，採用了這個領域大模型之前都不太關注微調技術。 這些結果的大部分來自對論文直接引用，並使用相同的數據劃分策略，有一些沒有 benchmark 的任務，研究團隊使用了 ESM-15B + finetuning 的策略作為 benchmark， 實際上，在所有任務中，研究團隊也使用過這種策略，發現 ESM2-15B/3B + finetuning 在不少任務可以直接達到 SOTA，但是，目前大部分的蛋白質大語言模型很少關注微調技術，更多的是把 PLMs 作為特徵提取器使用。
另一方面，為了表明 xTrimoPGLM 大模型的有效性，研究團隊也給出來和 ESM2-15B 以及 150M 的在相同訓練設置下的對比實驗
總的來說，xTrimoPGLM-100B 在 15 項任務中的 12 項上表現得比 ESM2-15B 更好。研究結果也揭示了一個規律：儘管其它的方法（比如 Ankh）在試圖尋找一種途徑，即在不依賴大規模語言模型的情況下，構建高效、低成本、有知識指導的蛋白質語言模型，但在模型的性能還是與模型規模密切相關，這表明，擴大模型規模可能是一個簡單而有效的方法，能夠在通用的蛋白質相關任務中提升模型的表現。這為未來對蛋白質預訓練模型進一步的研究提供了指導。 抗體藥物設計的兩項任務 為了進一步確定 xTrimoPGLM 框架的通用性，團隊把該框架應用在抗體蛋白預訓練上。考慮到訓練資源的限制和抗體數據的多樣性不足（大部分長度相似且有相似的框架區域），團隊沒有直接在 xTrimoPGLM-100B 上進行精調，而是構建了一個 12 億參數模型 xTrimoPGLM-Ab-1B，在包含 10 億抗體序列的 OAS 數據集對模型進行訓練。 考慮到 CDRs 是抗體最重要的部分，團隊對 40% 的樣本進行完整的 CDRs 掩碼處理，另外 40% 的樣本隨機掩碼處理，而剩下的 20% 則使用 MLM 目標。由於 [gMASK] 在抗體任務中的需求較少，所以沒採用該 loss。xTrimoPGLM-Ab-1B 先在通用蛋白序列上訓練 500B 的 token，隨後在 OAS 的數據上接著訓練 500B 的 token。一共使用 128 塊 Nvidia A100 80G GPU 卡進行混合精度訓練，大約需要 168 小時。 對於 antibody-based 的藥物設計，有兩項必不可少的任務就是序列自然度 (Naturalness) 以及抗體結構預測 (Antibody Structure Prediction), 下面分別介紹。 Zero-shot Naturalness 研究團隊使用了百圖生科濕實驗室獲得的蛋白質表達實驗數據集來評估各種模型的性能。具體來說，任何產生的純化蛋白質少於 10 mg/L 的樣本都被歸類為未表達，而產生超過 10 mg/L 的樣本被認為是成功合成的。第一個數據集（數據集 1）包括 601 個抗體序列，來自在 CHO 和 HEK293 細胞上進行的濕實驗。其中，成功表達的有 516 個。第二個數據集（數據集 2）包含了 98 個針對特定抗原的人類抗體序列，其中 90 個成功表達。評估採用 zero-shot 方式評估，不針對標籤微調，僅通過計算序列困惑度（PPL）和偽困惑度（PPPL）給序列打分。 結果顯示，在這兩個數據集中，xTrimoPGLM-Ab-1B 均優於其他基準模型。而且，進一步對 xTrimoPGLM-Ab-1B 進行精調，分別得到了 xTrimoPGLM-Ab-1B-GLM 和 xTrimoPGLM-Ab-1B-MLM 兩個模型。結果顯示，這兩個模型在數據集 2 上的 AUC 得分均有 0.02 的提升。
作者們還進行了消融研究，以證明隨機區域掩碼和 CDR 區域掩碼的重要性。實驗證明，同時使用這兩種目標在數據集 1 和數據集 2 上的表現均優於僅使用其中一種任務的模型，這突顯了組合使用這兩類目標的重要性。 抗體結構預測 這個任務目標是根據抗體的序列來預測其結構，實驗涵蓋了單鏈結構預測和複雜結構預測，即 VH-VL 複合物。 單鏈結構預測的數據集源自 2022 年 4 月 13 日之前的 RCSB Protein Data Bank (PDB)。該數據集包含 19k 個抗體鏈（VL 或 VH）。通過過濾，最終獲得了約 7.5k 個獨特的序列。另一個數據集，VH-VL 複合物，包含了大約 4.7k 個來自 PDB 的抗體。評估標準為根均方偏差（RMSD）和 TM-score。複雜結構預測還包括 DockQ 評估。 相比當前流行的結構預測模型（如 ESMFold, AlphaFold2），xTrimoPGLM-AbFold 做了以下改變：1) 去除了 MSA 和模板搜索模塊；2) 將下游 evoformer 模塊的數量從 48 減少到 1。
對於單鏈結構預測任務，研究團隊對比了 Alphafold2 以及四個基於 PLM 的模型：OmegaFold、ESMFold、IgFold 和 xTrimoAbFold。
結果如表所示，xTrimoPGLM-AbFold 在所有抗體結構預測相關的指標上都顯著優於其他模型，進一步說明，在預訓練模型對數據分布擬合的足夠好時，只需微調一個額外的 Evoformer 模塊以及不依賴 MSA 和模板的情況下，就能成為領先的抗體結構預測模型。 對於 VH-VL 複雜結構的預測，研究團隊比較了 ZDock、ClusPro、EquiDock、HDOCK 以及 AlphaFold-Multimer。
上表也展示了不同模型在 VH-VL 複合物性能上的表現。AlphaFold-Multimer 使用了 MSA 和模板信息，效果優於大多數結構預測演算法。而 xTrimoPGLM-AbFold 不使用任何 MSA 或模板信息，與 AlphaFold-Multimer 的表現相當，這說明 xTrimoPGLM-Ab-1B 已經學習了足夠豐富的抗體信息。更重要的是，其推理速度比 AlphaFold-Multimer 快了 6300 倍，而且比使用了 MSA 加速搜索策略的 AlphaFold-Multimer 快了 103 倍。在 AI 製藥引擎中，往往需要對生成的候選序列快速進行結構預測，以便算出 reward，進行下一輪的迭代，速度的提升決定了引擎的效率 此外，團隊將 Evoformer 模塊的數量增加到 16 個時，xTrimoPGLM-AbFold 在所有指標上都達到了最佳性能，同時速度還比原來的 AlphaFold-Multimer 快 2400 倍，比加速 MSA 搜索版的 AlphaFold-Multimer 快 40 倍。值得注意的是，當 Evoformer 模塊的數量從 1 增加到 16 時，效果只有少量的提升，這表明預訓練模型已經學習到了足夠的序列信息，可以精確地預測原子位置。 普通蛋白生成 在探索 xTrimoPGLM 生成自然功能序列的能力上，研究團隊生成了數千個序列並預測其對應的三維摺疊結構。研究團隊發現該模型能夠生成重要的二級結構，包括 alpha 螺旋和 beta 摺疊，這些都是更複雜三級結構的基礎。此外，模型生成的序列與自然序列相似性低，這為之後的藥物合成提供了更多的選擇。
但是，如何生成高質量序列還有著巨大的挑戰。首先，當處理超過 200 個氨基酸的序列時，模型往往生成大量的循環，而不是自然蛋白質般的結構。其次，模型在生成過程中經常會產生重複的問題，這可能源於模型傾向於選擇在局部最大化輸出概率的選項。對此團隊嘗試使用了 n-gram penalty 來減少生成重複序列的可能性，但研究團隊發現許多示例都展示出低複雜度的序列（例如，局部重複），預測的結構中包含長循環無序區域，推測是 n-gram 懲罰可能阻礙了模型生成正確序列的能力（下圖第一行）。在去除 n-gram penalty 後，模型能夠生成正常的結構（下圖第二行）。
抗體蛋白生成 團隊近一步展示了 xTrimoPGLM-Ab-1B 模型的生成能力，選擇了一個能與 SARS-CoV-2-WT 鏈接的重鏈抗體序列，並採用四種不同的策略重新設計了該序列的 CDR3 區域。因為 CDR3 區域在抗體或 T 細胞受體的結構中起著關鍵作用，並且具有顯著的變異性，對抗原識別的特異性起著重要的作用。以下是定義的四種策略： CDR3 短序列掩碼（CSM）：對 CDR3 部分區域進行掩碼重設計。 CDR3 全序列掩碼（CWM）：對 CDR3 全部區域進行掩碼重設計。 CDR3 隨機突變（CRM）：對 CDR3 區域內特定位點的隨機 3-6 個位置進行突變。 CDR3 隨機檢索（CRR）：使用來自 SARS-CoV-2 野生型庫中其他抗體的序列隨機替換 CDR3 區域。 團隊通過 xTrimoPGLM-Ab-1B 生成了一組 6,000 個抗體。研究團隊隨機選擇了六種抗體，並使用 xTrimoPGLM-AbFold 作為結構預測模型。CSM 和 CWM 策略能夠生成不同長度的序列，而不必進行突變或刪除。相比之下，由兩個並行基線 CRM 和 CRR 生成的序列顯示出相當大的無序性，無論是否存在少量突變或完全替換整個 CDR3 片段。研究團隊的分析進一步發現，編輯距離與生成抗體 CDR3 區域的結構之間存在關係。特別是，隨著編輯距離的增加，CDR3 區域的結構傾向於退化，即使是大的生成模型目前也仍然面臨限制。
結語 通過借鑒 NLP/CV 領域的想法，生物領域的預訓練模型近兩年雨後春筍般地冒出來，隨著模型計算力的提升和生物數據增長，我們期待更多未知的、驚人的發現出現在這個領域中。 儘管仍然存在不少需要繼續探索的地方，千億模型的誕生不僅標誌著最前沿的 AI 技術和生物學技術的融合，還意味著一個充滿無限可能的未來已經開啟。我們期待，這一重量級的模型引領製藥領域步入一個新的黃金時代，為人類健康和科學事業開創更加光明的未來。 © THE END
轉載請聯繫本公眾號獲得授權
「將蛋白質語言模型擴展到千億參數，深度解讀百圖生科、清華xTrimoPGLM模型」相關視頻
