4000萬蛋白結(jié)構(gòu)訓練，西湖大學開發(fā)基于結(jié)構(gòu)詞表的蛋白質(zhì)通用大模型，已開源
蛋白質(zhì)結(jié)構(gòu)相比于序列往往被認為更加具有信息量，因為其直接決定了蛋白質(zhì)的功能。而隨著AlphaFold2帶來的巨大突破，大量的預測結(jié)構(gòu)被發(fā)布出來供人研究使用。如何利用這些蛋白質(zhì)結(jié)構(gòu)來訓練強大且通用的表征模型是一個值得研究的方向。
西湖大學的研究人員利用Foldseek來處理蛋白質(zhì)結(jié)構(gòu)，將其編碼成一維的離散token，并與傳統(tǒng)的氨基酸進行結(jié)合，形成了結(jié)構(gòu)感知詞表（Structure-aware Vocabulary），以此將結(jié)構(gòu)信息嵌入到模型輸入中，增強模型的表征能力。
在預訓練上，論文使用了目前最多的蛋白質(zhì)結(jié)構(gòu)（identity過濾后4000萬），在64張A100上訓練了3個月，最終開源了具備650M參數(shù)量的模型SaProt（同時包括了35M的版本）。實驗結(jié)果表明SaProt各種蛋白質(zhì)任務上都要好于之前的序列和結(jié)構(gòu)模型。
研究《SaProt: Protein Language Modeling with Structure-aware Vocabulary》的預印版本，于 2024 年 3 月 21 日發(fā)布在 bioRxiv 預印平臺。
論文鏈接：https://www.biorxiv.org/content/10.1101/2023.10.01.560349v4
github鏈接：https://github.com/westlake-repl/SaProt
方法
本文利用Foldseek將蛋白質(zhì)進行編碼，生成了一維的3Di結(jié)構(gòu)序列（使用了Foldseek的結(jié)構(gòu)詞表，每種3Di token代表不同的局部結(jié)構(gòu)），這樣的結(jié)構(gòu)序列與氨基酸序列是等長的。
因此本文使用了一種簡單而有效的結(jié)構(gòu)嵌入方式：將結(jié)構(gòu)詞表和氨基酸詞表計算笛卡爾積（即兩兩組合），形成新的結(jié)構(gòu)感知詞表。這樣對于蛋白質(zhì)的每個位點，其氨基酸類型和對應的局部結(jié)構(gòu)都能組合成新詞表中的某個元素，從而讓模型同時考慮到蛋白質(zhì)的序列與結(jié)構(gòu)信息。
本文使用Bert架構(gòu)進行掩碼語言建模（Masked Language Modeling ）預訓練（關(guān)于訓練的更多細節(jié)可參考原論文）。
圖：結(jié)構(gòu)感知詞表
實驗
方法對比
一個可能令人疑惑的問題就是為什么需要這樣編碼結(jié)構(gòu)？論文展示了使用不同的結(jié)構(gòu)編碼方式進行預訓練的結(jié)果圖：
圖：不同結(jié)構(gòu)模型訓練的loss曲線圖
圖左和圖中是兩種經(jīng)典的蛋白質(zhì)結(jié)構(gòu)建模方式，即將結(jié)構(gòu)信息編碼成bias后添加到transformer的attention map中（如Evoformer，Uni-Mol），或者使用圖神經(jīng)網(wǎng)絡的方式建模蛋白質(zhì)的空間關(guān)系（如MIF，GearNet等）。
然而從loss圖中可以發(fā)現(xiàn)，當上述兩種建模方式在AF2結(jié)構(gòu)上使用MLM的訓練目標進行預訓練時，模型會非常迅速地過擬合（表現(xiàn)為在AF2預測結(jié)構(gòu)上預測loss非常低，但在PDB真實結(jié)構(gòu)上loss停滯甚至上升）。
作者推測這是由于AF2預測出來的蛋白質(zhì)結(jié)構(gòu)帶有一些隱藏的模式（patterns），由于前兩種方式是直接對蛋白質(zhì)的三維坐標進行建模，這些隱藏的pattern可能很輕易地就被模型識別出來，從而造成了信息泄露的問題，讓模型無需真正學習到蛋白質(zhì)的進化信息就能輕松地完成訓練目標。
而結(jié)構(gòu)感知詞表通過將蛋白質(zhì)結(jié)構(gòu)編碼成一維的結(jié)構(gòu)序列，在盡可能保留結(jié)構(gòu)模式的情況下忽略了精細的坐標數(shù)值，因此模型能夠有效地利用結(jié)構(gòu)信息而不受到隱藏pattern的影響。
Zero-shot測試
作者在蛋白質(zhì)突變數(shù)據(jù)集（ProteinGym）上和真實人類臨床疾病數(shù)據(jù)集（ClinVar）上測試了SaProt的zero-shot能力，結(jié)果如下：
圖：Zero-shot實驗結(jié)果
SaProt在兩個數(shù)據(jù)集上都超越了以往的所有結(jié)構(gòu)和序列模型，證明了其在zero-shot預測突變上具備優(yōu)異的能力。
監(jiān)督微調(diào)測試
本文還涵蓋了各種下游任務來測試模型表現(xiàn)，結(jié)果如下：
圖：下游任務fine-tune結(jié)果
SaProt在各個下游任務上都超越了以往的序列和結(jié)構(gòu)模型，展示出了其強大且通用的表征能力。
結(jié)構(gòu)信息測試
SaProt在4000萬的蛋白質(zhì)結(jié)構(gòu)上進行訓練，獲得了強大的表征能力。一個可能的疑問是如何確定SaProt學到了更多的結(jié)構(gòu)信息而不是模型被訓練得更好？
論文對SaProt和ESM-2在殘基接觸預測任務（Contact Prediction Task）上進行了測試。作者凍住了模型的backbone，只訓練一個線性分類層。實驗結(jié)果如下：
圖：Contact Prediction Task的結(jié)果
從結(jié)果可以看到，由于結(jié)構(gòu)token的嵌入，SaProt的表現(xiàn)大大超越了ESM-2，這表明SaProt蘊含了非常豐富的結(jié)構(gòu)信息，使其能夠在結(jié)構(gòu)預測任務上獲得十分優(yōu)異的結(jié)果。同時，論文在SCOPe數(shù)據(jù)庫上對alpha蛋白質(zhì)和beta蛋白質(zhì)進行了可視化，結(jié)果如下：
圖：在SCOPe數(shù)據(jù)庫上的Embedding可視化
SaProt的可視化結(jié)果非常清晰地將alpha蛋白質(zhì)和beta蛋白質(zhì)區(qū)分開來，而ESM-2的可視化結(jié)果卻將兩種蛋白質(zhì)混雜在一起，這說明了SaProt對結(jié)構(gòu)的變化有很強的感知能力。
不同結(jié)構(gòu)預測方法的比較
除了AF2，目前還存在許多其他的單序列結(jié)構(gòu)預測方法（如ESMFold），因此本文額外測試了其他方法預測出來的結(jié)構(gòu)對SaProt性能的作用。結(jié)果如下：
圖：不同結(jié)構(gòu)預測方法的fine-tune結(jié)果
從測試結(jié)果可以看出，雖然SaProt在AF2結(jié)構(gòu)上的表現(xiàn)最好（模型本身也是基于AF2結(jié)構(gòu)進行訓練的），但其他的結(jié)構(gòu)預測方法也能讓SaProt與ESM-2等模型性能相當。這意味著考慮到計算與時間成本，單序列結(jié)構(gòu)預測模型也能作為替代方法輸入到SaProt中。
局限
雖然SaProt經(jīng)過訓練展示出了優(yōu)異的性能，但依然還有一些可以改進的地方，例如：
Foldseek默認的結(jié)構(gòu)詞表大小只有20，如果有更加精準的結(jié)構(gòu)編碼模型，擴大結(jié)構(gòu)表征的詞表大小，是不是能進一步提升模型利用結(jié)構(gòu)的能力？
由于計算能力的限制，SaProt只在650M上完成了訓練。如果能夠繼續(xù)擴大模型規(guī)模，是否可以進一步地提升模型表現(xiàn)？
論文雖然已經(jīng)測試了很多的蛋白質(zhì)任務，但還有一些其他任務可以應用探索，例如蛋白質(zhì)序列設計（給定backbone預測氨基酸序列）等。
