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