論文鏈接：https://www.biorxiv.org/content/10.1101/2023.07.05.547496v3

自然語言處理（NLP）領(lǐng)域中的預(yù)訓(xùn)練模型相關(guān)進(jìn)展已經(jīng)被成功地用于學(xué)習(xí)蛋白質(zhì)序列中隱藏的生物信息?，F(xiàn)在大多數(shù)的蛋白質(zhì)預(yù)訓(xùn)練模型都受限于自動編碼或自回歸預(yù)訓(xùn)練目標(biāo)，這使得它們難以同時處理蛋白質(zhì)理解（例如，蛋白質(zhì)結(jié)構(gòu)預(yù)測）和生成任務(wù)（例如，藥物設(shè)計）。

這篇論文提出統(tǒng)一的蛋白質(zhì)語言模型，xTrimoPGLM，通過一個創(chuàng)新的預(yù)訓(xùn)練框架來同時處理這兩種類型的任務(wù)。xTrimoPGLM 主要技術(shù)貢獻(xiàn)是探索了這兩種類型目標(biāo)之間的兼容性以及共同優(yōu)化的可能性，并基于此訓(xùn)練了一個前所未有的 1000 億參數(shù)規(guī)模的蛋白質(zhì)語言模型，并消耗了 1 萬億 Tokens，模型 FLOPs 達(dá)到 6.2e+23，達(dá)到和 175B 參數(shù)的 GPT-3 一個量級。

在理解任務(wù)上，xTrimoPGLM 在多種蛋白質(zhì)理解任務(wù)（15 項(xiàng)任務(wù)中的 13 項(xiàng)任務(wù)）中顯著優(yōu)于其他先進(jìn)基線。在生成任務(wù)上，xTrimoPGLM 能夠生成與自然蛋白質(zhì)結(jié)構(gòu)類似的新蛋白質(zhì)序列。

此外，文章基于相同的框架額外訓(xùn)練了一個 12 億參數(shù)的抗體模型（xTrimoPGLM-Ab），其在預(yù)測抗體自然性和結(jié)構(gòu)方面取得了市面上最好的效果，并且顯示出比 AlphaFold2 更快的推理速度（數(shù)十倍到數(shù)千倍）。綜合來看，這些結(jié)果充分展示了 xTrimoPGLM 在理解和生成蛋白質(zhì)序列方面的強(qiáng)大能力和廣闊的應(yīng)用前景。

蛋白質(zhì)理解和生成的統(tǒng)一

不同類型的蛋白質(zhì)相關(guān)任務(wù)需要蛋白質(zhì)語言模型（Protein Language Model，PLM）提供各異的輸出。具體的，蛋白質(zhì)理解任務(wù)，如二級結(jié)構(gòu)預(yù)測等，需要 PLM 提供精確的氨基酸和序列級別的表示；而蛋白質(zhì)設(shè)計任務(wù)，如抗體或酶的設(shè)計，依賴于 PLM 的生成能力。然而，當(dāng)前的 PLM 因?yàn)槠鋯我坏念A(yù)訓(xùn)練框架的限制，大多只能處理一種類型的任務(wù)。

事實(shí)上，蛋白質(zhì)的理解和生成都反映了蛋白質(zhì)數(shù)據(jù)的分布信息，Meta 之前使用 ESM（基于 Masked Language Model 的蛋白質(zhì)大模型）做生成的工作也證實(shí)了這一點(diǎn)，指出蛋白質(zhì)理解預(yù)訓(xùn)練模型可以通過一些采樣策略進(jìn)行蛋白質(zhì)設(shè)計。這進(jìn)一步支撐了這兩種看似不同的任務(wù)的統(tǒng)一性，如果能夠使用同一個訓(xùn)練框架去處理這兩種任務(wù)，將會進(jìn)一步增強(qiáng)模型對蛋白質(zhì)數(shù)據(jù)的擬合能力。

雖然 NLP 領(lǐng)域生成式模型（例如 UL2R, GPT) 已經(jīng)成為主流范式，通過把數(shù)據(jù)標(biāo)簽映射到整個文本空間，結(jié)合指令微調(diào)來生成各式各樣的任務(wù)的答案，但 PLM 還無法實(shí)現(xiàn)這一點(diǎn)。實(shí)際上，蛋白質(zhì)的應(yīng)用仍然依賴于將表示與下游任務(wù)特定標(biāo)簽（如結(jié)構(gòu)預(yù)測的 3D 坐標(biāo)）之間的橋接，這在很大程度上依賴于 BERT 樣式的訓(xùn)練來處理蛋白質(zhì)理解任務(wù)。因此，需要同時進(jìn)行這兩種訓(xùn)練目標(biāo)。

本文介紹的 xTrimo 蛋白質(zhì)通用語言模型（xTrimoPGLM）預(yù)訓(xùn)練框架，巧妙地統(tǒng)一了兩種類型的預(yù)訓(xùn)練任務(wù)，使模型能同時處理多種與蛋白質(zhì)相關(guān)的任務(wù)。研究團(tuán)隊(duì)通過全面的實(shí)驗(yàn)評估了 xTrimoPGLM 框架的有效性。在蛋白質(zhì)理解任務(wù)的情況下，xTrimoPGLM-100B 在多種評估中表現(xiàn)出色，涵蓋了蛋白質(zhì)結(jié)構(gòu)、功能、交互和可開發(fā)性等領(lǐng)域的 15 項(xiàng)任務(wù)。

"Scaling Law" 是衡量大型語言模型的重要原則，模型的性能應(yīng)隨著模型參數(shù)大小、數(shù)據(jù)量、計算量按比例指數(shù)增加而線性增長。實(shí)際上，研究團(tuán)隊(duì)在下游任務(wù)上的實(shí)驗(yàn)結(jié)果驗(yàn)證了這一定律，證明了大型模型在處理復(fù)雜任務(wù)時的必要性。

如下圖所示，性能改善與預(yù)訓(xùn)練計算量增加之間的關(guān)系。以 Meta 的 ESM-2 為參考，隨著蛋白質(zhì)語言模型（PLM）的計算量呈指數(shù)增長，蛋白質(zhì)的下游性能仍然會線性增長（每個大類有 3-4 個任務(wù)，數(shù)值表示這些任務(wù)的平均值）。

圖片

在蛋白質(zhì)生成任務(wù)中，xTrimoPGLM-100B 展示了生成不同長度和序列的新蛋白質(zhì)序列的能力，這是通過調(diào)整生成超參數(shù)實(shí)現(xiàn)的。值得注意的是，當(dāng)與自然蛋白質(zhì)進(jìn)行比較時，xTrimoPGLM-100B 展示了生成結(jié)構(gòu)相似但序列相異的新蛋白質(zhì)序列的能力。這再次驗(yàn)證了大型模型對于復(fù)雜任務(wù)的重要性，進(jìn)一步證實(shí)了遵循 "Scaling Law" 進(jìn)行模型設(shè)計的決定是正確的。

同時，研究團(tuán)隊(duì)還開發(fā)了一種當(dāng)前更具實(shí)用意義的具有 12 億參數(shù)的抗體 PLM，即 xTrimoPGLM-Ab-1B。這種模型在 OAS 抗體數(shù)據(jù)庫上進(jìn)行 fine-tuning，處理了超過 1 萬億個 token。它在抗體的自然性和結(jié)構(gòu)預(yù)測任務(wù)上達(dá)到了目前最優(yōu)秀的性能。由于不依賴于外部庫的檢索和多序列對齊（Multiple Sequence Alignment)，所以結(jié)構(gòu)預(yù)測在速度上比 Alphafold2 模型提升成百上千倍，這對于基于抗體藥物發(fā)現(xiàn) AI 制藥公司至關(guān)重要。

混合訓(xùn)練

原始的 GLM 模型利用兩種不同的預(yù)訓(xùn)練目標(biāo)來提升其生成能力：1）跨度生成（Span Generation，簡稱 sMask），用于恢復(fù)句子中的短空白；2）長文本生成（Long-text Generation，簡稱 gMASK），用于在提供前綴上下文的基礎(chǔ)上生成隨機(jī)長度的序列。為了進(jìn)一步提升 xTrimoPGLM 的理解能力，團(tuán)隊(duì)在 prefix 區(qū)域引入了被用作理解目標(biāo)的 Masked Language Model（MLM，即 [MASK]）。這樣的設(shè)計確保了 xTrimoPGLM 能生成精確的殘基級和序列級表示。

當(dāng)使用 [MASK] 標(biāo)識符時，xTrimoPGLM 的功能類似于 BERT。相反，當(dāng)使用 [sMASK] 或 [gMASK] 時，xTrimoPGLM 的行為類似于 PrefixLM 或 GPT?？偟膩碚f，xTrimoPGLM-100B 的預(yù)訓(xùn)練階段可以分為兩個階段。首先，利用 MLM 進(jìn)行預(yù)訓(xùn)練以增強(qiáng)其表示能力，主要目標(biāo)是快速減少損失水平。第二階段，使用結(jié)合 MLM 和 GLM 損失的統(tǒng)一目標(biāo)進(jìn)行訓(xùn)練，以提升理解和生成能力。

圖片

NLP 領(lǐng)域大量探索了統(tǒng)一的預(yù)訓(xùn)練模式，但大多還是采樣了同樣的訓(xùn)練模式（自回歸或自編碼）。為了滿足統(tǒng)一的蛋白質(zhì)預(yù)訓(xùn)練模型的需求，需要將 BERT 樣式的目標(biāo)引入到預(yù)訓(xùn)練語言模型中，以增強(qiáng)模型的表示能力，同時也需要引入 GPT 樣式的目標(biāo)，以確保模型的生成能力。在最開始研究團(tuán)隊(duì)使用 Probing 策略探索 Contact Map Prediction 的任務(wù)時，發(fā)現(xiàn)僅僅依靠基于下一個詞預(yù)測的生成式語言模型，效果會有大幅度的下降。

兼容性實(shí)驗(yàn)：在實(shí)證分析中，研究團(tuán)隊(duì)在 xTrimoPGLM-150m 模型上探究了同時優(yōu)化兩個不同目標(biāo)的可行性。結(jié)果如下圖 (a)(b) 所示，即使兩種預(yù)訓(xùn)練目標(biāo)看似沖突，MLM 損失和 GLM 損失也可以同時優(yōu)化，反之亦然，即統(tǒng)一的訓(xùn)練可以很快的適配到 MLM 或者 GLM 上，并不會影響兩者的收斂情況。

研究團(tuán)隊(duì)還探究了 MLM 與 GLM 兩種目標(biāo)是否能相互加速收斂，結(jié)果如圖 (c)(d)。 

• MLM-adapted GLM：接著 MLM 預(yù)訓(xùn)練后的模型，繼續(xù)訓(xùn)練 GLM 目標(biāo)函數(shù)；
• GLM-adapted MLM：接著 GLM 預(yù)訓(xùn)練后的模型，繼續(xù)訓(xùn)練 MLM 目標(biāo)函數(shù)；

總的說來，基于預(yù)訓(xùn)練后的模型中繼續(xù)訓(xùn)練的模型，與從頭開始訓(xùn)練的模型相比，其收斂速度明顯加快。這些實(shí)驗(yàn)觀察到：在蛋白質(zhì)數(shù)據(jù)分布并不依賴于特定的訓(xùn)練模式，從而縮小了自編碼 PLMs（如 ESM）和自回歸 PLMs（如 ProGen2）之間的差距，為 100B 訓(xùn)練流程提供了支持。

圖片

訓(xùn)練穩(wěn)定性

訓(xùn)練穩(wěn)定性是成功訓(xùn)練 100B 規(guī)模大型語言模型的決定性因素。xTrimoPGLM 從 GLM-130B 的實(shí)現(xiàn)中借鑒了一些想法并解決了許多不穩(wěn)定訓(xùn)練的問題。然而，xTrimoPGLM-100B 在從訓(xùn)練的第一階段過渡到第二階段時仍然會遇到災(zāi)難性的訓(xùn)練崩潰（小規(guī)模的模型（10B 規(guī)模）中并未觀察到），即使一開始只將 1% 的 GLM 損失加入預(yù)訓(xùn)練也可能觸發(fā)這些崩潰。

下圖可以看到，如果直接給 GLM 分配一個比例 ，在訓(xùn)練初期，grad norm 都會出現(xiàn) spike（橙，藍(lán)，綠線）。 

圖片

為了緩解這個問題，研究團(tuán)隊(duì)提出了一種平滑過渡策略。主要分為兩個階段。在第一階段，主要目標(biāo)是逐步提高 GLM 損失的比例，達(dá)到預(yù)期的數(shù)量。具體來說，給定一個期望的 GLM 損失比例 R，按照線性增長，以 K 步從 0 增加到 R。在這個階段應(yīng)該將學(xué)習(xí)率保持在極低的水平。完成過渡后，學(xué)習(xí)率可以按照預(yù)定義的腳本在幾百個 steps 內(nèi)逐漸回升至原來的水平 （紫線）。實(shí)際上，最后的 xTrimoPGLM-100B 訓(xùn)練運(yùn)行只在過渡階段經(jīng)歷了損失分歧情況，但是由于硬件故障多次失敗，導(dǎo)致經(jīng)常性的更換節(jié)點(diǎn)和重啟。

訓(xùn)練數(shù)據(jù)

為了能夠盡可能的映射整個蛋白世界， xTrimoPGLM-100B 的預(yù)訓(xùn)練模型的訓(xùn)練數(shù)據(jù)集整合自兩個廣泛的數(shù)據(jù)源：Uniref90 和 ColAbFoldDB。

結(jié)合這兩個數(shù)據(jù)源，預(yù)訓(xùn)練模型數(shù)據(jù)集充分利用了這兩個數(shù)據(jù)源的優(yōu)勢，既有廣泛的生物分類覆蓋，又有多樣的環(huán)境生態(tài)位蛋白質(zhì)序列，全面而詳盡地映射了生物世界中的蛋白質(zhì)資源。

圖片

超參數(shù)配置

xTrimoPGLM-100B 模型的訓(xùn)練過程復(fù)雜且耗費(fèi)資源，團(tuán)隊(duì)耗費(fèi)了 160 天的時間，開發(fā)團(tuán)隊(duì)使用了 96 臺 DGX-A100 GPU 服務(wù)器（每臺服務(wù)器擁有 8×40G 的 GPU）以混合精度（FP16）進(jìn)行訓(xùn)練，消耗了 1 萬億的 tokens，由于大多數(shù)現(xiàn)有的大型語言模型在訓(xùn)練上存在嚴(yán)重不足，所以開發(fā)團(tuán)隊(duì)現(xiàn)在仍然在繼續(xù)訓(xùn)練 xTrimoPGLM-100B 模型，以處理盡可能多的 tokens。

模型采用了 3D 并行策略，基于 DeepSpeed 進(jìn)行了 4 路張量并行、8 路流水線并行，并且采用了 Zero Stage 1 進(jìn)行訓(xùn)練。模型具有 72 個 Transformer 層，80 個 head，以及 10,240 維和 31,744 個 FNN 的維度。使用了 DeepNorm 初始化 Post-LN，并且采用了 Embedding Layer Gradient Shrink (EGS) 穩(wěn)定訓(xùn)練，以及使用了 2D ROPE 位置編碼技術(shù)。為了提升訓(xùn)練效率，每個樣本由多個蛋白拼接再一起，使用 < eos > 區(qū)分開來，包含固定的 2,048 的序列長度。最終單卡的 TFLOPs 在 120-135 之間，68 examples/sec,  如果是 80G 的 A100，經(jīng)過減少重算可到 92 examples/sec。下表展示了大部分超參數(shù)的配置。

圖片

團(tuán)隊(duì)也對比了當(dāng)前比較熱門的預(yù)訓(xùn)練模型的 FLOPs，目前生物領(lǐng)域不同模型的結(jié)構(gòu)也存在比較大的區(qū)別， 所以在 FLOPs 的計算上，團(tuán)隊(duì)考慮盡可能詳細(xì)統(tǒng)計，包括 query、 key、value 的轉(zhuǎn)換，Attention Matrix 的計算，注意力之后線性變換，MLP 中的變換，以及最后到 vocab 維度的映射，以及語言模型頭（如果存在的話）中的線性轉(zhuǎn)換等，可以看到，xTrimoPGLM-100B 高出其他模型一到兩個量級。

圖片

十五種蛋白質(zhì)下游任務(wù)評估

為了全面評估 xTrimoPGLM-100B 模型，研究團(tuán)隊(duì)對 4 個領(lǐng)域內(nèi)，15 個下游的蛋白質(zhì)相關(guān)任務(wù)進(jìn)行了基準(zhǔn)測試。測試結(jié)果表明，xTrimoPGLM-100B 在蛋白質(zhì)結(jié)構(gòu)、蛋白質(zhì)可開發(fā)能力、蛋白質(zhì)相互作用和蛋白質(zhì)功能等四個主要類別的任務(wù)中，都表現(xiàn)出了顯著的優(yōu)勢。在這些任務(wù)中，xTrimoPGLM-100B 模型與微調(diào)技術(shù)的組合取得了優(yōu)異的成績，大部分超越了當(dāng)前最先進(jìn)的方法，從而推動這個領(lǐng)域的進(jìn)步。

需要強(qiáng)調(diào)的是，下表對比主要從任務(wù)的角度進(jìn)行，而不是一個完全公平的對比，因?yàn)?xTrimoPGLM-100B 在取得這些結(jié)果時，采用了這個領(lǐng)域大模型之前都不太關(guān)注微調(diào)技術(shù)。這些結(jié)果的大部分來自對論文直接引用，并使用相同的數(shù)據(jù)劃分策略，有一些沒有 benchmark 的任務(wù)，研究團(tuán)隊(duì)使用了 ESM-15B + finetuning 的策略作為 benchmark， 實(shí)際上，在所有任務(wù)中，研究團(tuán)隊(duì)也使用過這種策略，發(fā)現(xiàn) ESM2-15B/3B +  finetuning 在不少任務(wù)可以直接達(dá)到 SOTA，但是，目前大部分的蛋白質(zhì)大語言模型很少關(guān)注微調(diào)技術(shù)，更多的是把 PLMs 作為特征提取器使用。

圖片

另一方面，為了表明 xTrimoPGLM 大模型的有效性，研究團(tuán)隊(duì)也給出來和 ESM2-15B 以及 150M 的在相同訓(xùn)練設(shè)置下的對比實(shí)驗(yàn)，使用相對較小的 ESM2-150M 模型作為指標(biāo)，主要來理解各種下游蛋白質(zhì)相關(guān)任務(wù)的難度程度。這些任務(wù)包括基于 feature-based 的 Probing 和 聯(lián)合大模型參數(shù)的 Finetuning，xTrimoPGLM-100B 在大多數(shù)蛋白質(zhì)相關(guān)任務(wù)中仍然展現(xiàn)出了優(yōu)勢。

圖片

總的來說，xTrimoPGLM-100B 在 15 項(xiàng)任務(wù)中的 12 項(xiàng)上表現(xiàn)得比 ESM2-15B 更好。研究結(jié)果也揭示了一個規(guī)律：盡管其它的方法（比如 Ankh）在試圖尋找一種途徑，即在不依賴大規(guī)模語言模型的情況下，構(gòu)建高效、低成本、有知識指導(dǎo)的蛋白質(zhì)語言模型，但在模型的性能還是與模型規(guī)模密切相關(guān)，這表明，擴(kuò)大模型規(guī)?？赡苁且粋€簡單而有效的方法，能夠在通用的蛋白質(zhì)相關(guān)任務(wù)中提升模型的表現(xiàn)。這為未來對蛋白質(zhì)預(yù)訓(xùn)練模型進(jìn)一步的研究提供了指導(dǎo)。

抗體藥物設(shè)計的兩項(xiàng)任務(wù)

為了進(jìn)一步確定 xTrimoPGLM 框架的通用性，團(tuán)隊(duì)把該框架應(yīng)用在抗體蛋白預(yù)訓(xùn)練上?？紤]到訓(xùn)練資源的限制和抗體數(shù)據(jù)的多樣性不足（大部分長度相似且有相似的框架區(qū)域），團(tuán)隊(duì)沒有直接在 xTrimoPGLM-100B 上進(jìn)行精調(diào)，而是構(gòu)建了一個 12 億參數(shù)模型 xTrimoPGLM-Ab-1B，在包含 10 億抗體序列的 OAS 數(shù)據(jù)集對模型進(jìn)行訓(xùn)練。

考慮到 CDRs 是抗體最重要的部分，團(tuán)隊(duì)對 40% 的樣本進(jìn)行完整的 CDRs 掩碼處理，另外 40% 的樣本隨機(jī)掩碼處理，而剩下的 20% 則使用 MLM 目標(biāo)。由于 [gMASK] 在抗體任務(wù)中的需求較少，所以沒采用該 loss。xTrimoPGLM-Ab-1B 先在通用蛋白序列上訓(xùn)練 500B 的 token，隨后在 OAS 的數(shù)據(jù)上接著訓(xùn)練 500B 的 token。一共使用 128 塊 Nvidia A100 80G GPU 卡進(jìn)行混合精度訓(xùn)練，大約需要 168 小時。

對于 antibody-based 的藥物設(shè)計，有兩項(xiàng)必不可少的任務(wù)就是序列自然度 (Naturalness) 以及抗體結(jié)構(gòu)預(yù)測 (Antibody Structure Prediction), 下面分別介紹。

Zero-shot Naturalness

研究團(tuán)隊(duì)使用了百圖生科濕實(shí)驗(yàn)室獲得的蛋白質(zhì)表達(dá)實(shí)驗(yàn)數(shù)據(jù)集來評估各種模型的性能。具體來說，任何產(chǎn)生的純化蛋白質(zhì)少于 10 mg/L 的樣本都被歸類為未表達(dá)，而產(chǎn)生超過 10 mg/L 的樣本被認(rèn)為是成功合成的。第一個數(shù)據(jù)集（數(shù)據(jù)集 1）包括 601 個抗體序列，來自在 CHO 和 HEK293 細(xì)胞上進(jìn)行的濕實(shí)驗(yàn)。其中，成功表達(dá)的有 516 個。第二個數(shù)據(jù)集（數(shù)據(jù)集 2）包含了 98 個針對特定抗原的人類抗體序列，其中 90 個成功表達(dá)。評估采用 zero-shot 方式評估，不針對標(biāo)簽微調(diào)，僅通過計算序列困惑度（PPL）和偽困惑度（PPPL）給序列打分。

結(jié)果顯示，在這兩個數(shù)據(jù)集中，xTrimoPGLM-Ab-1B 均優(yōu)于其他基準(zhǔn)模型。而且，進(jìn)一步對 xTrimoPGLM-Ab-1B 進(jìn)行精調(diào)，分別得到了 xTrimoPGLM-Ab-1B-GLM 和 xTrimoPGLM-Ab-1B-MLM 兩個模型。結(jié)果顯示，這兩個模型在數(shù)據(jù)集 2 上的 AUC 得分均有 0.02 的提升。

圖片

作者們還進(jìn)行了消融研究，以證明隨機(jī)區(qū)域掩碼和 CDR 區(qū)域掩碼的重要性。實(shí)驗(yàn)證明，同時使用這兩種目標(biāo)在數(shù)據(jù)集 1 和數(shù)據(jù)集 2 上的表現(xiàn)均優(yōu)于僅使用其中一種任務(wù)的模型，這突顯了組合使用這兩類目標(biāo)的重要性。

抗體結(jié)構(gòu)預(yù)測

這個任務(wù)目標(biāo)是根據(jù)抗體的序列來預(yù)測其結(jié)構(gòu)，實(shí)驗(yàn)涵蓋了單鏈結(jié)構(gòu)預(yù)測和復(fù)雜結(jié)構(gòu)預(yù)測，即 VH-VL 復(fù)合物。

單鏈結(jié)構(gòu)預(yù)測的數(shù)據(jù)集源自 2022 年 4 月 13 日之前的 RCSB Protein Data Bank (PDB)。該數(shù)據(jù)集包含 19k 個抗體鏈（VL 或 VH）。通過過濾，最終獲得了約 7.5k 個獨(dú)特的序列。另一個數(shù)據(jù)集，VH-VL 復(fù)合物，包含了大約 4.7k 個來自 PDB 的抗體。評估標(biāo)準(zhǔn)為根均方偏差（RMSD）和 TM-score。復(fù)雜結(jié)構(gòu)預(yù)測還包括 DockQ 評估。

相比當(dāng)前流行的結(jié)構(gòu)預(yù)測模型（如 ESMFold, AlphaFold2），xTrimoPGLM-AbFold 做了以下改變：1) 去除了 MSA 和模板搜索模塊；2) 將下游 evoformer 模塊的數(shù)量從 48 減少到 1。

圖片

對于單鏈結(jié)構(gòu)預(yù)測任務(wù)，研究團(tuán)隊(duì)對比了 Alphafold2 以及四個基于 PLM 的模型：OmegaFold、ESMFold、IgFold 和 xTrimoAbFold。

圖片

結(jié)果如表所示，xTrimoPGLM-AbFold 在所有抗體結(jié)構(gòu)預(yù)測相關(guān)的指標(biāo)上都顯著優(yōu)于其他模型，進(jìn)一步說明，在預(yù)訓(xùn)練模型對數(shù)據(jù)分布擬合的足夠好時，只需微調(diào)一個額外的 Evoformer 模塊以及不依賴 MSA 和模板的情況下，就能成為領(lǐng)先的抗體結(jié)構(gòu)預(yù)測模型。

對于 VH-VL 復(fù)雜結(jié)構(gòu)的預(yù)測，研究團(tuán)隊(duì)比較了 ZDock、ClusPro、EquiDock、HDOCK 以及 AlphaFold-Multimer。

圖片

上表也展示了不同模型在 VH-VL 復(fù)合物性能上的表現(xiàn)。AlphaFold-Multimer 使用了 MSA 和模板信息，效果優(yōu)于大多數(shù)結(jié)構(gòu)預(yù)測算法。而 xTrimoPGLM-AbFold 不使用任何 MSA 或模板信息，與 AlphaFold-Multimer 的表現(xiàn)相當(dāng)，這說明 xTrimoPGLM-Ab-1B 已經(jīng)學(xué)習(xí)了足夠豐富的抗體信息。更重要的是，其推理速度比 AlphaFold-Multimer 快了 6300 倍，而且比使用了 MSA 加速搜索策略的 AlphaFold-Multimer 快了 103 倍。在 AI 制藥引擎中，往往需要對生成的候選序列快速進(jìn)行結(jié)構(gòu)預(yù)測，以便算出 reward，進(jìn)行下一輪的迭代，速度的提升決定了引擎的效率。

此外，團(tuán)隊(duì)將 Evoformer 模塊的數(shù)量增加到 16 個時，xTrimoPGLM-AbFold 在所有指標(biāo)上都達(dá)到了最佳性能，同時速度還比原來的 AlphaFold-Multimer 快 2400 倍，比加速 MSA 搜索版的 AlphaFold-Multimer 快 40 倍。值得注意的是，當(dāng) Evoformer 模塊的數(shù)量從 1 增加到 16 時，效果只有少量的提升，這表明預(yù)訓(xùn)練模型已經(jīng)學(xué)習(xí)到了足夠的序列信息，可以精確地預(yù)測原子位置。

普通蛋白生成

在探索 xTrimoPGLM 生成自然功能序列的能力上，研究團(tuán)隊(duì)生成了數(shù)千個序列并預(yù)測其對應(yīng)的三維折疊結(jié)構(gòu)。研究團(tuán)隊(duì)發(fā)現(xiàn)該模型能夠生成重要的二級結(jié)構(gòu)，包括 alpha 螺旋和 beta 折疊，這些都是更復(fù)雜三級結(jié)構(gòu)的基礎(chǔ)。此外，模型生成的序列與自然序列相似性低，這為之后的藥物合成提供了更多的選擇。

圖片

但是，如何生成高質(zhì)量序列還有著巨大的挑戰(zhàn)。首先，當(dāng)處理超過 200 個氨基酸的序列時，模型往往生成大量的循環(huán)，而不是自然蛋白質(zhì)般的結(jié)構(gòu)。其次，模型在生成過程中經(jīng)常會產(chǎn)生重復(fù)的問題，這可能源于模型傾向于選擇在局部最大化輸出概率的選項(xiàng)。對此團(tuán)隊(duì)嘗試使用了 n-gram penalty 來減少生成重復(fù)序列的可能性，但研究團(tuán)隊(duì)發(fā)現(xiàn)許多示例都展示出低復(fù)雜度的序列（例如，局部重復(fù)），預(yù)測的結(jié)構(gòu)中包含長循環(huán)無序區(qū)域，推測是 n-gram 懲罰可能阻礙了模型生成正確序列的能力（下圖第一行）。在去除 n-gram penalty 后，模型能夠生成正常的結(jié)構(gòu)（下圖第二行）。

圖片

抗體蛋白生成

團(tuán)隊(duì)近一步展示了 xTrimoPGLM-Ab-1B 模型的生成能力，選擇了一個能與 SARS-CoV-2-WT 鏈接的重鏈抗體序列，并采用四種不同的策略重新設(shè)計了該序列的 CDR3 區(qū)域。因?yàn)?CDR3 區(qū)域在抗體或 T 細(xì)胞受體的結(jié)構(gòu)中起著關(guān)鍵作用，并且具有顯著的變異性，對抗原識別的特異性起著重要的作用。以下是定義的四種策略：

• CDR3 短序列掩碼（CSM）：對 CDR3 部分區(qū)域進(jìn)行掩碼重設(shè)計。
• CDR3 全序列掩碼（CWM）：對 CDR3 全部區(qū)域進(jìn)行掩碼重設(shè)計。
• CDR3 隨機(jī)突變（CRM）：對 CDR3 區(qū)域內(nèi)特定位點(diǎn)的隨機(jī) 3-6 個位置進(jìn)行突變。
• CDR3 隨機(jī)檢索（CRR）：使用來自 SARS-CoV-2 野生型庫中其他抗體的序列隨機(jī)替換 CDR3 區(qū)域。

團(tuán)隊(duì)通過 xTrimoPGLM-Ab-1B 生成了一組 6,000 個抗體。研究團(tuán)隊(duì)隨機(jī)選擇了六種抗體，并使用 xTrimoPGLM-AbFold 作為結(jié)構(gòu)預(yù)測模型。CSM 和 CWM 策略能夠生成不同長度的序列，而不必進(jìn)行突變或刪除。相比之下，由兩個并行基線 CRM 和 CRR 生成的序列顯示出相當(dāng)大的無序性，無論是否存在少量突變或完全替換整個 CDR3 片段。研究團(tuán)隊(duì)的分析進(jìn)一步發(fā)現(xiàn)，編輯距離與生成抗體 CDR3 區(qū)域的結(jié)構(gòu)之間存在關(guān)系。特別是，隨著編輯距離的增加，CDR3 區(qū)域的結(jié)構(gòu)傾向于退化，即使是大的生成模型目前也仍然面臨限制。

圖片

結(jié)語

通過借鑒 NLP/CV 領(lǐng)域的想法，生物領(lǐng)域的預(yù)訓(xùn)練模型近兩年雨后春筍般地冒出來，隨著模型計算力的提升和生物數(shù)據(jù)增長，我們期待更多未知的、驚人的發(fā)現(xiàn)出現(xiàn)在這個領(lǐng)域中。

盡管仍然存在不少需要繼續(xù)探索的地方，千億模型的誕生不僅標(biāo)志著最前沿的 AI 技術(shù)和生物學(xué)技術(shù)的融合，還意味著一個充滿無限可能的未來已經(jīng)開啟。我們期待，這一重量級的模型引領(lǐng)制藥領(lǐng)域步入一個新的黃金時代，為人類健康和科學(xué)事業(yè)開創(chuàng)更加光明的未來。