蛋白質(zhì)結(jié)構(gòu)相比于序列往往被認(rèn)為更加具有信息量，因?yàn)槠渲苯記Q定了蛋白質(zhì)的功能。而隨著AlphaFold2帶來的巨大突破，大量的預(yù)測(cè)結(jié)構(gòu)被發(fā)布出來供人研究使用。如何利用這些蛋白質(zhì)結(jié)構(gòu)來訓(xùn)練強(qiáng)大且通用的表征模型是一個(gè)值得研究的方向。

西湖大學(xué)的研究人員利用Foldseek來處理蛋白質(zhì)結(jié)構(gòu)，將其編碼成一維的離散token，并與傳統(tǒng)的氨基酸進(jìn)行結(jié)合，形成了結(jié)構(gòu)感知詞表（Structure-aware Vocabulary），以此將結(jié)構(gòu)信息嵌入到模型輸入中，增強(qiáng)模型的表征能力。

在預(yù)訓(xùn)練上，論文使用了目前最多的蛋白質(zhì)結(jié)構(gòu)（identity過濾后4000萬），在64張A100上訓(xùn)練了3個(gè)月，最終開源了具備650M參數(shù)量的模型SaProt（同時(shí)包括了35M的版本）。實(shí)驗(yàn)結(jié)果表明SaProt各種蛋白質(zhì)任務(wù)上都要好于之前的序列和結(jié)構(gòu)模型。

研究《SaProt: Protein Language Modeling with Structure-aware Vocabulary》的預(yù)印版本，于 2024 年 3 月 21 日發(fā)布在 bioRxiv 預(yù)印平臺(tái)。

論文鏈接：https://www.biorxiv.org/content/10.1101/2023.10.01.560349v4

github鏈接：https://github.com/westlake-repl/SaProt

方法

本文利用Foldseek將蛋白質(zhì)進(jìn)行編碼，生成了一維的3Di結(jié)構(gòu)序列（使用了Foldseek的結(jié)構(gòu)詞表，每種3Di token代表不同的局部結(jié)構(gòu)），這樣的結(jié)構(gòu)序列與氨基酸序列是等長(zhǎng)的。

因此本文使用了一種簡(jiǎn)單而有效的結(jié)構(gòu)嵌入方式：將結(jié)構(gòu)詞表和氨基酸詞表計(jì)算笛卡爾積（即兩兩組合），形成新的結(jié)構(gòu)感知詞表。這樣對(duì)于蛋白質(zhì)的每個(gè)位點(diǎn)，其氨基酸類型和對(duì)應(yīng)的局部結(jié)構(gòu)都能組合成新詞表中的某個(gè)元素，從而讓模型同時(shí)考慮到蛋白質(zhì)的序列與結(jié)構(gòu)信息。

本文使用Bert架構(gòu)進(jìn)行掩碼語言建模（Masked Language Modeling ）預(yù)訓(xùn)練（關(guān)于訓(xùn)練的更多細(xì)節(jié)可參考原論文）。

圖：結(jié)構(gòu)感知詞表

實(shí)驗(yàn)

方法對(duì)比

一個(gè)可能令人疑惑的問題就是為什么需要這樣編碼結(jié)構(gòu)？論文展示了使用不同的結(jié)構(gòu)編碼方式進(jìn)行預(yù)訓(xùn)練的結(jié)果圖：

圖：不同結(jié)構(gòu)模型訓(xùn)練的loss曲線圖

圖左和圖中是兩種經(jīng)典的蛋白質(zhì)結(jié)構(gòu)建模方式，即將結(jié)構(gòu)信息編碼成bias后添加到transformer的attention map中（如Evoformer，Uni-Mol），或者使用圖神經(jīng)網(wǎng)絡(luò)的方式建模蛋白質(zhì)的空間關(guān)系（如MIF，GearNet等）。

然而從loss圖中可以發(fā)現(xiàn)，當(dāng)上述兩種建模方式在AF2結(jié)構(gòu)上使用MLM的訓(xùn)練目標(biāo)進(jìn)行預(yù)訓(xùn)練時(shí)，模型會(huì)非常迅速地過擬合（表現(xiàn)為在AF2預(yù)測(cè)結(jié)構(gòu)上預(yù)測(cè)loss非常低，但在PDB真實(shí)結(jié)構(gòu)上loss停滯甚至上升）。

作者推測(cè)這是由于AF2預(yù)測(cè)出來的蛋白質(zhì)結(jié)構(gòu)帶有一些隱藏的模式（patterns），由于前兩種方式是直接對(duì)蛋白質(zhì)的三維坐標(biāo)進(jìn)行建模，這些隱藏的pattern可能很輕易地就被模型識(shí)別出來，從而造成了信息泄露的問題，讓模型無需真正學(xué)習(xí)到蛋白質(zhì)的進(jìn)化信息就能輕松地完成訓(xùn)練目標(biāo)。

而結(jié)構(gòu)感知詞表通過將蛋白質(zhì)結(jié)構(gòu)編碼成一維的結(jié)構(gòu)序列，在盡可能保留結(jié)構(gòu)模式的情況下忽略了精細(xì)的坐標(biāo)數(shù)值，因此模型能夠有效地利用結(jié)構(gòu)信息而不受到隱藏pattern的影響。

Zero-shot測(cè)試

作者在蛋白質(zhì)突變數(shù)據(jù)集（ProteinGym）上和真實(shí)人類臨床疾病數(shù)據(jù)集（ClinVar）上測(cè)試了SaProt的zero-shot能力，結(jié)果如下：

圖：Zero-shot實(shí)驗(yàn)結(jié)果

SaProt在兩個(gè)數(shù)據(jù)集上都超越了以往的所有結(jié)構(gòu)和序列模型，證明了其在zero-shot預(yù)測(cè)突變上具備優(yōu)異的能力。

監(jiān)督微調(diào)測(cè)試

本文還涵蓋了各種下游任務(wù)來測(cè)試模型表現(xiàn)，結(jié)果如下：

圖：下游任務(wù)fine-tune結(jié)果

SaProt在各個(gè)下游任務(wù)上都超越了以往的序列和結(jié)構(gòu)模型，展示出了其強(qiáng)大且通用的表征能力。

結(jié)構(gòu)信息測(cè)試

SaProt在4000萬的蛋白質(zhì)結(jié)構(gòu)上進(jìn)行訓(xùn)練，獲得了強(qiáng)大的表征能力。一個(gè)可能的疑問是如何確定SaProt學(xué)到了更多的結(jié)構(gòu)信息而不是模型被訓(xùn)練得更好？

論文對(duì)SaProt和ESM-2在殘基接觸預(yù)測(cè)任務(wù)（Contact Prediction Task）上進(jìn)行了測(cè)試。作者凍住了模型的backbone，只訓(xùn)練一個(gè)線性分類層。實(shí)驗(yàn)結(jié)果如下：

圖：Contact Prediction Task的結(jié)果

從結(jié)果可以看到，由于結(jié)構(gòu)token的嵌入，SaProt的表現(xiàn)大大超越了ESM-2，這表明SaProt蘊(yùn)含了非常豐富的結(jié)構(gòu)信息，使其能夠在結(jié)構(gòu)預(yù)測(cè)任務(wù)上獲得十分優(yōu)異的結(jié)果。同時(shí)，論文在SCOPe數(shù)據(jù)庫(kù)上對(duì)alpha蛋白質(zhì)和beta蛋白質(zhì)進(jìn)行了可視化，結(jié)果如下：

圖：在SCOPe數(shù)據(jù)庫(kù)上的Embedding可視化

SaProt的可視化結(jié)果非常清晰地將alpha蛋白質(zhì)和beta蛋白質(zhì)區(qū)分開來，而ESM-2的可視化結(jié)果卻將兩種蛋白質(zhì)混雜在一起，這說明了SaProt對(duì)結(jié)構(gòu)的變化有很強(qiáng)的感知能力。

不同結(jié)構(gòu)預(yù)測(cè)方法的比較

除了AF2，目前還存在許多其他的單序列結(jié)構(gòu)預(yù)測(cè)方法（如ESMFold），因此本文額外測(cè)試了其他方法預(yù)測(cè)出來的結(jié)構(gòu)對(duì)SaProt性能的作用。結(jié)果如下：

圖：不同結(jié)構(gòu)預(yù)測(cè)方法的fine-tune結(jié)果

從測(cè)試結(jié)果可以看出，雖然SaProt在AF2結(jié)構(gòu)上的表現(xiàn)最好（模型本身也是基于AF2結(jié)構(gòu)進(jìn)行訓(xùn)練的），但其他的結(jié)構(gòu)預(yù)測(cè)方法也能讓SaProt與ESM-2等模型性能相當(dāng)。這意味著考慮到計(jì)算與時(shí)間成本，單序列結(jié)構(gòu)預(yù)測(cè)模型也能作為替代方法輸入到SaProt中。

局限

雖然SaProt經(jīng)過訓(xùn)練展示出了優(yōu)異的性能，但依然還有一些可以改進(jìn)的地方，例如：

Foldseek默認(rèn)的結(jié)構(gòu)詞表大小只有20，如果有更加精準(zhǔn)的結(jié)構(gòu)編碼模型，擴(kuò)大結(jié)構(gòu)表征的詞表大小，是不是能進(jìn)一步提升模型利用結(jié)構(gòu)的能力？

由于計(jì)算能力的限制，SaProt只在650M上完成了訓(xùn)練。如果能夠繼續(xù)擴(kuò)大模型規(guī)模，是否可以進(jìn)一步地提升模型表現(xiàn)？

論文雖然已經(jīng)測(cè)試了很多的蛋白質(zhì)任務(wù)，但還有一些其他任務(wù)可以應(yīng)用探索，例如蛋白質(zhì)序列設(shè)計(jì)（給定backbone預(yù)測(cè)氨基酸序列）等。