蛋白質(zhì)結(jié)構(gòu)相比于序列往往被認(rèn)為更加具有信息量，因?yàn)槠渲苯記Q定了蛋白質(zhì)的功能。

而隨著AlphaFold2帶來(lái)的巨大突破，大量的預(yù)測(cè)結(jié)構(gòu)被發(fā)布出來(lái)供人研究使用，如何利用這些蛋白質(zhì)結(jié)構(gòu)來(lái)訓(xùn)練強(qiáng)大且通用的表征模型是一個(gè)值得研究的方向。

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

論文地址：https://www.biorxiv.org/content/10.1101/2023.10.01.560349v4

Github地址：https://github.com/westlake-repl/SaProt

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

ProteinGym：https://proteingym.org/benchmarks

方法

該研究利用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)，其氨基酸類(lèi)型和對(duì)應(yīng)的局部結(jié)構(gòu)都能組合成新詞表中的某個(gè)元素，從而讓模型同時(shí)考慮到蛋白質(zhì)的序列與結(jié)構(gòu)信息。

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

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

實(shí)驗(yàn)

方法對(duì)比

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

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

圖左和圖中是兩種經(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ì)非常迅速地過(guò)擬合（表現(xiàn)為在AF2預(yù)測(cè)結(jié)構(gòu)上預(yù)測(cè)loss非常低，但在PDB真實(shí)結(jié)構(gòu)上loss停滯甚至上升）。

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

而結(jié)構(gòu)感知詞表通過(guò)將蛋白質(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í)人類(lèi)臨床疾病數(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ù)來(lái)測(cè)試模型表現(xiàn)，結(jié)果如下：

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

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

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

SaProt在4000萬(wàn)的蛋白質(zhì)結(jié)構(gòu)上進(jìn)行訓(xùn)練，獲得了強(qiáng)大的表征能力。一個(gè)可能的疑問(wèn)是如何確定SaProt學(xué)到了更多的結(jié)構(gòu)信息而不是模型被訓(xùn)練得更好？論文對(duì)SaProt和ESM-2在殘基接觸預(yù)測(cè)任務(wù)（Contact Prediction Task）上進(jìn)行了測(cè)試。作者凍住了模型的backbone，只訓(xùn)練一個(gè)線(xiàn)性分類(lèi)層。

實(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ū)分開(kāi)來(lái)，而ESM-2的可視化結(jié)果卻將兩種蛋白質(zhì)混雜在一起，這說(shuō)明了SaProt對(duì)結(jié)構(gòu)的變化有很強(qiáng)的感知能力。

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

除了AF2，目前還存在許多其他的單序列結(jié)構(gòu)預(yù)測(cè)方法（如ESMFold），因此本文額外測(cè)試了其他方法預(yù)測(cè)出來(lái)的結(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)過(guò)訓(xùn)練展示出了優(yōu)異的性能，但依然還有一些可以改進(jìn)的地方，例如：

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

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

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