論文鏈接：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ì)序列中隱藏的生物信息。現(xiàn)在大多數(shù)的蛋白質(zhì)預(yù)訓(xùn)練模型都受限于自動(dòng)編碼或自回歸預(yù)訓(xùn)練目標(biāo)，這使得它們難以同時(shí)處理蛋白質(zhì)理解（例如，蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)）和生成任務(wù)（例如，藥物設(shè)計(jì)）。

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

在理解任務(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)練了一個(gè) 12 億參數(shù)的抗體模型（xTrimoPGLM-Ab），其在預(yù)測(cè)抗體自然性和結(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ù)，如二級(jí)結(jié)構(gòu)預(yù)測(cè)等，需要 PLM 提供精確的氨基酸和序列級(jí)別的表示；而蛋白質(zhì)設(shè)計(jì)任務(wù)，如抗體或酶的設(shè)計(jì)，依賴于 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ì)。這進(jìn)一步支撐了這兩種看似不同的任務(wù)的統(tǒng)一性，如果能夠使用同一個(gè)訓(xùn)練框架去處理這兩種任務(wù)，將會(huì)進(jìn)一步增強(qiáng)模型對(duì)蛋白質(zhì)數(shù)據(jù)的擬合能力。

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

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

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

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

圖片

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

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

混合訓(xùn)練

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

當(dāng)使用 [MASK] 標(biāo)識(shí)符時(shí)，xTrimoPGLM 的功能類似于 BERT。相反，當(dāng)使用 [sMASK] 或 [gMASK] 時(shí)，xTrimoPGLM 的行為類似于 PrefixLM 或 GPT。總的來說，xTrimoPGLM-100B 的預(yù)訓(xùn)練階段可以分為兩個(gè)階段。首先，利用 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)模型的表示能力，同時(shí)也需要引入 GPT 樣式的目標(biāo)，以確保模型的生成能力。在最開始研究團(tuán)隊(duì)使用 Probing 策略探索 Contact Map Prediction 的任務(wù)時(shí)，發(fā)現(xiàn)僅僅依靠基于下一個(gè)詞預(yù)測(cè)的生成式語言模型，效果會(huì)有大幅度的下降。

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

研究團(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)練的第一階段過渡到第二階段時(shí)仍然會(huì)遇到災(zāi)難性的訓(xùn)練崩潰（小規(guī)模的模型（10B 規(guī)模）中并未觀察到），即使一開始只將 1% 的 GLM 損失加入預(yù)訓(xùn)練也可能觸發(fā)這些崩潰。

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

圖片

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

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

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

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

圖片

超參數(shù)配置

xTrimoPGLM-100B 模型的訓(xùn)練過程復(fù)雜且耗費(fèi)資源，團(tuán)隊(duì)耗費(fèi)了 160 天的時(shí)間，開發(fā)團(tuán)隊(duì)使用了 96 臺(tái) DGX-A100 GPU 服務(wù)器（每臺(tái)服務(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 個(gè) Transformer 層，80 個(gè) head，以及 10,240 維和 31,744 個(gè) FNN 的維度。使用了 DeepNorm 初始化 Post-LN，并且采用了 Embedding Layer Gradient Shrink (EGS) 穩(wěn)定訓(xùn)練，以及使用了 2D ROPE 位置編碼技術(shù)。為了提升訓(xùn)練效率，每個(gè)樣本由多個(gè)蛋白拼接再一起，使用 < eos > 區(qū)分開來，包含固定的 2,048 的序列長(zhǎng)度。最終單卡的 TFLOPs 在 120-135 之間，68 examples/sec,  如果是 80G 的 A100，經(jīng)過減少重算可到 92 examples/sec。下表展示了大部分超參數(shù)的配置。

圖片

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

圖片

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

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

需要強(qiáng)調(diào)的是，下表對(duì)比主要從任務(wù)的角度進(jìn)行，而不是一個(gè)完全公平的對(duì)比，因?yàn)?xTrimoPGLM-100B 在取得這些結(jié)果時(shí)，采用了這個(gè)領(lǐng)域大模型之前都不太關(guān)注微調(diào)技術(shù)。這些結(jié)果的大部分來自對(duì)論文直接引用，并使用相同的數(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è)置下的對(duì)比實(shí)驗(yàn)，使用相對(duì)較小的 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)勢(shì)。

圖片

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

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

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

考慮到 CDRs 是抗體最重要的部分，團(tuán)隊(duì)對(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 小時(shí)。

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

Zero-shot Naturalness

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

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

圖片

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

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

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

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

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

圖片

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

圖片

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

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

圖片

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

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

普通蛋白生成

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

圖片

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

圖片

抗體蛋白生成

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

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

團(tuán)隊(duì)通過 xTrimoPGLM-Ab-1B 生成了一組 6,000 個(gè)抗體。研究團(tuán)隊(duì)隨機(jī)選擇了六種抗體，并使用 xTrimoPGLM-AbFold 作為結(jié)構(gòu)預(yù)測(cè)模型。CSM 和 CWM 策略能夠生成不同長(zhǎng)度的序列，而不必進(jìn)行突變或刪除。相比之下，由兩個(gè)并行基線 CRM 和 CRR 生成的序列顯示出相當(dāng)大的無序性，無論是否存在少量突變或完全替換整個(gè) 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)練模型近兩年雨后春筍般地冒出來，隨著模型計(jì)算力的提升和生物數(shù)據(jù)增長(zhǎng)，我們期待更多未知的、驚人的發(fā)現(xiàn)出現(xiàn)在這個(gè)領(lǐng)域中。

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