PLM重大突破！上海交大與上海AI Lab最新成果入選NeurIPS 24，ProSST有效整合蛋白

蛋白質(zhì)作為生命體的關(guān)鍵分子，其序列決定結(jié)構(gòu)，結(jié)構(gòu)決定功能，蛋白質(zhì)的功能與其三維結(jié)構(gòu)緊密相關(guān)。數(shù)十年來(lái)，科學(xué)家們利用 X 射線(xiàn)晶體學(xué)、核磁共振等技術(shù)，解析了數(shù)千種蛋白質(zhì)結(jié)構(gòu)，為人們理解蛋白質(zhì)功能提供了重要線(xiàn)索。然而，面對(duì)數(shù)百萬(wàn)級(jí)別的蛋白質(zhì)數(shù)量，解析所有蛋白質(zhì)結(jié)構(gòu)的任務(wù)顯得異常艱巨。
**受自然語(yǔ)言處理領(lǐng)域中預(yù)訓(xùn)練語(yǔ)言模型的啟發(fā)，預(yù)訓(xùn)練的蛋白質(zhì)語(yǔ)言模型 (PLMs) 應(yīng)運(yùn)而生。**通過(guò)在海量未標(biāo)記的蛋白質(zhì)序列數(shù)據(jù)上進(jìn)行學(xué)習(xí)，PLM 能夠捕捉到蛋白質(zhì)序列的復(fù)雜模式和相互作用，這為蛋白質(zhì)功能的預(yù)測(cè)、結(jié)構(gòu)分析以及蛋白質(zhì)-蛋白質(zhì)相互作用的識(shí)別帶來(lái)了革命性的進(jìn)步。

然而，大多數(shù) PLM 主要關(guān)注蛋白質(zhì)序列建模，忽略了結(jié)構(gòu)信息的重要性，這主要是因?yàn)榻Y(jié)構(gòu)數(shù)據(jù)的缺乏。隨著 AlphaFold 和 RoseTTAFold 等技術(shù)的出現(xiàn)，蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)的準(zhǔn)確性大幅提升，研究人員也開(kāi)始探索如何有效地將蛋白質(zhì)結(jié)構(gòu)信息整合到 PLM 中，從而訓(xùn)練大規(guī)模結(jié)構(gòu)感知的預(yù)訓(xùn)練語(yǔ)言模型。
比如，上海交通大學(xué)自然科學(xué)研究院/物理天文學(xué)院/張江高研院/藥學(xué)院洪亮教授課題組，上海交大助理研究員周冰心，聯(lián)合上海人工智能實(shí)驗(yàn)室青年研究員談攀，近期成功研發(fā)了一款具備結(jié)構(gòu)感知能力的預(yù)訓(xùn)練蛋白質(zhì)語(yǔ)言模型——ProSST。

具體來(lái)說(shuō)，該模型在包含 1,880 萬(wàn)蛋白質(zhì)結(jié)構(gòu)的大型數(shù)據(jù)集上進(jìn)行預(yù)訓(xùn)練，將蛋白質(zhì)結(jié)構(gòu)轉(zhuǎn)化為結(jié)構(gòu)化 token 序列，并與氨基酸序列一同輸入至 Transformer 模型中。通過(guò)采用解耦注意力 (disentangled attention) 機(jī)制，ProSST 能夠?qū)⑦@兩類(lèi)信息有效融合，從而在熱穩(wěn)定性預(yù)測(cè)、金屬離子結(jié)合預(yù)測(cè)、蛋白質(zhì)定位預(yù)測(cè)、GO 注釋預(yù)測(cè)等監(jiān)督學(xué)習(xí)任務(wù)中顯著超越現(xiàn)有模型。

該研究以「ProSST: Protein Language Modeling with Quantized Structure and Disentangled Attention」為題，已入選 NeurIPS 2024。

研究亮點(diǎn)：

* 該研究提出了一種蛋白質(zhì)結(jié)構(gòu)量化器，它能夠?qū)⒌鞍踪|(zhì)結(jié)構(gòu)轉(zhuǎn)換成一系列離散的結(jié)構(gòu)元，這些離散化的結(jié)構(gòu)元可有效表征蛋白質(zhì)中殘基的局部結(jié)構(gòu)信息

* 該研究提出了一種解耦注意力機(jī)制，用于學(xué)習(xí)蛋白質(zhì)氨基酸序列和三維結(jié)構(gòu)之間的關(guān)系，從而促進(jìn)結(jié)構(gòu)離散化序列與氨基酸序列的高效信息整合

* 與 ESM 系列、SaProt 等其他蛋白質(zhì)預(yù)訓(xùn)練大模型相比， ProSST 的參數(shù)量只有 110M，遠(yuǎn)小于經(jīng)典 ESM 系列的 650M，但 ProSST 在各種蛋白質(zhì)下游任務(wù)中幾乎都表現(xiàn)出最佳性能，體現(xiàn)了 ProSST 模型架構(gòu)設(shè)計(jì)的優(yōu)越性

* 在目前最大規(guī)模零樣本突變效應(yīng)預(yù)測(cè)的 ProteinGym Benchmark 上，ProSST 排名第一；在最新的 ProteinGym 上，是第一個(gè)把 zero-shot 突變性能預(yù)測(cè) Spearman 相關(guān)性做到 0.5 以上的開(kāi)源模型

基于業(yè)內(nèi)主流無(wú)監(jiān)督預(yù)訓(xùn)練數(shù)據(jù)集，覆蓋 1,880 萬(wàn)個(gè)蛋白質(zhì)結(jié)構(gòu)

為了實(shí)現(xiàn) ProSST 的無(wú)監(jiān)督預(yù)訓(xùn)練，研究團(tuán)隊(duì)主要使用了以下數(shù)據(jù)集：

* AlphaFoldDB 數(shù)據(jù)集：從超過(guò) 2.14 億個(gè)蛋白質(zhì)結(jié)構(gòu)中選取了縮減 90% 的版本 (reduced version)，共 1,880 萬(wàn)個(gè)結(jié)構(gòu)，其中隨機(jī)抽取 10 萬(wàn)個(gè)結(jié)構(gòu)作為驗(yàn)證集，用于監(jiān)測(cè)和調(diào)整訓(xùn)練階段的困惑度。

* CATH43-S40 數(shù)據(jù)集：包含 31,885 個(gè)經(jīng)過(guò) 40% 序列相似性去重的蛋白質(zhì)晶體結(jié)構(gòu)域，移除缺少關(guān)鍵原子（例如 Cα 和 N）的結(jié)構(gòu)后，剩下 31,270 條記錄，從中隨機(jī)抽取 200 個(gè)結(jié)構(gòu)作為驗(yàn)證集，用于監(jiān)控和優(yōu)化模型性能。

* CATH43-S40 局部結(jié)構(gòu)數(shù)據(jù)集：由 CATH43-S40 數(shù)據(jù)集中提取的局部結(jié)構(gòu)組成，通過(guò)構(gòu)建星形圖方法提取出 4,735,677 個(gè)局部結(jié)構(gòu)，用于結(jié)構(gòu)編碼器的嵌入表征和結(jié)構(gòu)碼本的聚類(lèi)分析。

* ProteinGYM 基準(zhǔn)數(shù)據(jù)集：用于評(píng)估 ProSST 在零樣本突變效應(yīng)預(yù)測(cè)方面的能力，包含 217 個(gè)實(shí)驗(yàn)分析，每個(gè)分析都包含蛋白質(zhì)的序列和結(jié)構(gòu)信息，特別關(guān)注 66 個(gè)聚焦于熱穩(wěn)定性的數(shù)據(jù)集，采用斯皮爾曼系數(shù)、Top-recall 和 NDCG 作為性能評(píng)估指標(biāo)。

ProSST：具有結(jié)構(gòu)感知能力的 PLM，包含兩個(gè)關(guān)鍵模塊

該研究開(kāi)發(fā)的 ProSST (Protein Sequence-Structure Transformer) 是一個(gè)具有結(jié)構(gòu)感知能力的預(yù)訓(xùn)練蛋白質(zhì)語(yǔ)言模型。如下圖所示，ProSST 主要由兩個(gè)模塊組成：結(jié)構(gòu)量化 (structure quantization) 模塊和具有序列-結(jié)構(gòu)解耦注意力 (sequence-structure disentangled attention) 的 Transformer 模型。

ProSST 的模型架構(gòu)

結(jié)構(gòu)量化模塊：將蛋白質(zhì)結(jié)構(gòu)序列化及量化為一系列結(jié)構(gòu)元

結(jié)構(gòu)量化模塊的目標(biāo)是將蛋白質(zhì)中殘基的局部結(jié)構(gòu)轉(zhuǎn)換為離散標(biāo)記。最初，局部結(jié)構(gòu) (local structure) 被預(yù)訓(xùn)練的結(jié)構(gòu)編碼器編碼成了一個(gè)密集向量。隨后，一個(gè)預(yù)訓(xùn)練的 k 均值聚類(lèi)模型 (k-means clustering model) 根據(jù)編碼向量為局部結(jié)構(gòu)分配一個(gè)類(lèi)別標(biāo)簽 (category label)。最后，將類(lèi)別標(biāo)簽作為結(jié)構(gòu)元 (structure token) 分配給殘基。

* 相對(duì)于蛋白質(zhì)整體結(jié)構(gòu)，局部結(jié)構(gòu)描述的顆粒度更細(xì)

結(jié)構(gòu)量化的流程

具體而言，該研究使用幾何向量感知器 (GVP) 作為局部結(jié)構(gòu)編碼器 (local structure encoder)，如下圖 A 所示，該研究將 GVP 與包含位置感知的多層感知器 (MLP) 的解碼器集成，形成了一個(gè)自編碼器模型。整個(gè)模型使用去噪預(yù)訓(xùn)練目標(biāo)蛋白進(jìn)行訓(xùn)練，在 C.A.T.H 數(shù)據(jù)集上訓(xùn)練后，研究人員僅使用編碼器的平均池化輸出作為結(jié)構(gòu)的最終表征。

結(jié)構(gòu)編碼器的訓(xùn)練

緊接著，如下圖 B 所示，該研究的結(jié)構(gòu)編碼器 (local structure encoder) 將表征蛋白質(zhì)結(jié)構(gòu)的密集向量量化為離散標(biāo)記。為此，研究人員使用結(jié)構(gòu)編碼器 GVP 將 C.A.T.H 數(shù)據(jù)集中所有殘基的局部結(jié)構(gòu) (local structures) 嵌入到連續(xù)的潛在空間中，然后應(yīng)用 k 均值 (k-means) 算法在這個(gè)潛在空間中識(shí)別 K 個(gè)質(zhì)心 (centroids)，這些質(zhì)心構(gòu)成了結(jié)構(gòu)碼本 (structure codebook)。

局部結(jié)構(gòu)聚類(lèi)和標(biāo)記

最后，對(duì)于蛋白質(zhì)序列中位置 i 的殘基，該研究首先基于其局部結(jié)構(gòu)構(gòu)建圖 Gi，然后使用結(jié)構(gòu)編碼器 GVP 將其嵌入到連續(xù)向量 (vector) ri 中?？傮w而言，如下圖 C 所示，整個(gè)蛋白質(zhì)結(jié)構(gòu)可以序列化并量化 (serialized and quantized) 為一系列結(jié)構(gòu)元 (structure tokens)。

將蛋白質(zhì)結(jié)構(gòu)轉(zhuǎn)化為結(jié)構(gòu)元序列

序列-結(jié)構(gòu)解耦注意力：使模型學(xué)習(xí)殘基與殘基、殘基與結(jié)構(gòu)之間的關(guān)系

該研究受到了 DeBerta 模型的啟發(fā)，旨在通過(guò)解耦注意力的方式來(lái)學(xué)習(xí)殘基序列（氨基酸序列）和結(jié)構(gòu)序列，以及相對(duì)位置三者間的相互關(guān)系，從而使得模型能夠處理蛋白質(zhì)序列和結(jié)構(gòu)信息，并通過(guò)解耦的方式來(lái)提高模型的性能和穩(wěn)定性。

具體來(lái)說(shuō)，對(duì)于蛋白質(zhì)一級(jí)序列中第 i 個(gè)殘基，可以通過(guò) 3 項(xiàng)來(lái)表示：Ri 表示氨基酸序列 token 的編碼，Si 代表氨基酸的局部結(jié)構(gòu) token 編碼，而 Pi|j 表示第 i 個(gè)殘基在位置 j 的 token 編碼。如下圖所示，該研究的序列-結(jié)構(gòu)解耦注意力機(jī)制包括殘基對(duì)殘基 (R to R)、殘基對(duì)結(jié)構(gòu) (R to S)、殘基對(duì)位置 (R to P)、結(jié)構(gòu)對(duì)殘基 (S to R)、位置對(duì)殘基 (P to R) 這 5 種類(lèi)型，使得模型能夠更細(xì)致地捕捉蛋白質(zhì)序列和結(jié)構(gòu)之間的復(fù)雜關(guān)系。

ProSST 的模型架構(gòu)

ProSST 性能全面領(lǐng)先，結(jié)構(gòu)信息的納入大幅提高模型表征能力

為了驗(yàn)證 ProSST 在零樣本突變效應(yīng)預(yù)測(cè) (zero-shot mutant effective prediction) 的有效性，該研究將其與多種頂尖模型進(jìn)行了比較，涵蓋序列模型 (sequence-based models)、結(jié)構(gòu)-序列模型 (structure-sequence model)、逆向折疊模型 (inverse folding models)、進(jìn)化模型 (evolutionary models) 和集成模型 (ensemble models)。

如下表所示， 在 ProteinGYM 的基準(zhǔn)評(píng)測(cè)下，ProSST 的表現(xiàn)優(yōu)于所有對(duì)比模型，并在穩(wěn)定性 (stability) 上達(dá)到最佳。此外，ProSST (-structure) 表現(xiàn)與其他序列模型相當(dāng)，這證實(shí)了 ProSST 的性能提升主要?dú)w功于其對(duì)結(jié)構(gòu)信息的有效整合。
* ProSST (-structure) 沒(méi)有納入結(jié)構(gòu)信息模塊

ProSST 與其他模型的零樣本突變預(yù)測(cè)性能比較

對(duì)于監(jiān)督學(xué)習(xí)，該研究選擇了熱穩(wěn)定性預(yù)測(cè) (Thermostability)、金屬離子結(jié)合預(yù)測(cè) (Metal Ion Binding)、蛋白質(zhì)定位預(yù)測(cè) (DeepLoc) 以及 GO 注釋預(yù)測(cè) (MF/BP/CC) 四大蛋白質(zhì)下游任務(wù)，并且將 ProSST 與 ESM-2、ESM-1b、SaProt、MIF-ST、GearNet 等其他蛋白質(zhì)語(yǔ)言模型進(jìn)行比較。結(jié)果如下表 2 所示，ProSST 在所有模型中獲得最佳結(jié)果，在所有 6 種設(shè)置中獲得了 5 個(gè)第一名和 1 個(gè)第二名。

下游任務(wù)的監(jiān)督微調(diào)比較

蛋白質(zhì)語(yǔ)言模型：連接大數(shù)據(jù)與生命科學(xué)的橋梁

自 ChatGPT 等大語(yǔ)言模型發(fā)布以來(lái)，基于大規(guī)模蛋白質(zhì)序列的預(yù)訓(xùn)練模型 (PLM) 已成為生命科學(xué)領(lǐng)域的熱門(mén)研究。目前，PLM 研究主要分為兩個(gè)方向：

* 檢索增強(qiáng)型 PLM：這類(lèi)模型會(huì)在訓(xùn)練或預(yù)測(cè)階段整合多序列比對(duì) (MSA) 信息，如 MSATransformer 和 Tranception，以提升預(yù)測(cè)性能。

* 多模態(tài) PLM：與僅使用序列信息的模型不同，多模態(tài) PLM 整合蛋白質(zhì)結(jié)構(gòu)等額外信息，如本文所述的 ProSST 模型，通過(guò)結(jié)構(gòu) token 序列與氨基酸序列的融合，即可增強(qiáng)模型的表征能力。

在檢索增強(qiáng)型 PLM 方面，今年 4 月，復(fù)旦大學(xué)等機(jī)構(gòu)的研究團(tuán)隊(duì)推出了 PLMSearch，這是一種基于序列輸入的同源蛋白質(zhì)搜索方法。該研究可利用預(yù)訓(xùn)練的蛋白質(zhì)語(yǔ)言模型獲取深層表征，并預(yù)測(cè)結(jié)構(gòu)相似度。相關(guān)研究已發(fā)表于 Nature Communication。

在多模態(tài) PLM 方面，浙江大學(xué)陳華鈞教授團(tuán)隊(duì)前不久提出了一種針對(duì)蛋白質(zhì)優(yōu)化的全新去噪蛋白質(zhì)語(yǔ)言模型 (DePLM)，該模型通過(guò)優(yōu)化進(jìn)化信息，可提高蛋白質(zhì)優(yōu)化任務(wù)的性能，相關(guān)成果已成功入選頂會(huì) NeurIPS 24。

隨著這些突破性研究的不斷涌現(xiàn)，PLM 正逐漸成為探索生命科學(xué)未知領(lǐng)域的強(qiáng)大工具，在蛋白質(zhì)功能預(yù)測(cè)、相互作用預(yù)測(cè)及表型關(guān)聯(lián)預(yù)測(cè)等領(lǐng)域具有巨大潛力，有望為疾病治療和改善人類(lèi)生活提供新思路。