作為細胞內無數生化反應的驅動力，蛋白質在細胞微觀世界中扮演著建筑師和工程師的角色，不僅催化著生命活動，更是構筑、維系生物體形態與功能的基礎構件。正是蛋白質之間的互動、協同作用，支撐起了生命的宏偉藍圖。
然而，蛋白質的結構復雜多變，傳統的實驗方法在解析蛋白質結構時既耗時又費力——蛋白質語言模型 (PLMs) 應運而生，利用深度學習技術，通過分析大量的蛋白質序列數據，學習蛋白質的生物化學規律和共進化模式，在蛋白質結構預測、適應性預測和蛋白質設計等領域取得了顯著成就，極大地推動了蛋白質工程的發展。
盡管 PLMs 在殘基尺度上取得了巨大成功，但在提供原子級信息方面的能力卻受到了限制。針對于此，清華大學智能產業研究院副研究員周浩聯合北京大學、南京大學和水木分子團隊，提出了一種多尺度的蛋白質語言模型 ESM-AA (ESM All Atom)， 通過設計殘基展開、多尺度位置編碼等訓練機制，拓展出了處理原子尺度信息的能力。
ESM-AA 在靶點-配體結合等任務的性能顯著提升，超越目前 SOTA 蛋白語言模型，如 ESM-2，也超越了目前的 SOTA 分子表示學習模型 Uni-Mol 等。相關研究已經以「ESM All-Atom: Multi-scale Protein Language Model for Unified Molecular Modeling」為題，發表于機器學習頂級會議 ICML 上。
論文地址：
https://icml.cc/virtual/2024/poster/35119
開源項目「awesome-ai4s」匯集了百余篇 AI4S 論文解讀，并提供海量數據集與工具：
https://github.com/hyperai/awesome-ai4s
數據集：構建了蛋白質和分子數據的混合數據集
在預訓練任務中，該研究使用了一個包含蛋白質和分子數據的組合數據集，其中包含原子坐標等結構信息。
對于蛋白質數據集，該研究使用了 AlphaFold DB ，其中包含了 800 萬條高置信度的 AlphaFold2 預測的蛋白質序列和結構。
對于分子數據集，該研究使用了由 ETKDG 和 MMFF 分子力場生成的數據，包含 1,900 萬個分子和 2.09 億種構型。
在訓練 ESM-AA 時，研究人員首先將一個蛋白質數據集 Dp 和一個分子數據集 Dm 混合在一起作為最終的數據集，即 D=Dp∪Dm 。對于來自 Dm 的分子而言，由于僅由原子組成，其代碼轉換序列 X? 就是所有原子 ā 的有序集合，并且沒有任何殘基，即 R?=?。值得注意的是，因為在預訓練中使用了分子數據，所以 ESM-AA 既可以接受蛋白質作為輸入，也可以接受分子作為輸入。
ESM-AA 模型構建：多尺度預訓練與編碼，實現統一分子建模
受多語言代碼切換方法的啟發，ESM-AA 在進行預測與蛋白質設計任務時，首先會隨機解壓縮部分殘基，從而生成多尺度代碼切換蛋白質序列，隨后通過精心設計的多尺度位置編碼對這些序列進行訓練，并且已經在殘基和原子尺度上證明了其有效性。
當處理蛋白質分子任務時，即涉及蛋白質和小分子的任務，ESM-AA 不需要任何額外模型輔助，可以充分發揮出預訓練模型的能力。
該研究的多尺度預訓練框架由多尺度掩碼語言建模 (masked language model, MLM) 和成對距離恢復 (pairwise distance recovery) 組成。
具體而言，在殘基尺度上，一個蛋白質 X 可以被看作是一個由 L 個殘基組成的序列，即 X = (r1,…,ri,…,rL)。每個殘基 ri 都是由 N 個原子 A 組成 Ai={a1i,…,aNi}。為了構建代碼切換蛋白序列 X?，該研究通過隨機選擇一組殘基，并將其對應的原子插入到 X 中，從而實現了一個解壓縮的過程。在這一過程中，研究人員將被解壓的原子按照順序排列，最后在將原子集合 Ai 插入到 X 中（即解壓殘基 ri）后，即可得到一個代碼切換序列 X?。
隨后，研究人員對代碼切換序列 X? 進行掩碼語言建模。
首先，隨機遮擋 X? 中的一部分原子或殘基，讓模型使用周圍上下文預測原始原子或殘基。然后，研究人員使用對偶距離恢復 (PDR) 作為另一個預訓練任務。即通過在坐標中添加噪聲來破壞原子尺度的結構信息，并使用被破壞的原子間距離信息作為模型輸入，要求模型恢復這些原子之間的準確歐幾里得距離。
考慮到跨越不同殘基的長距離結構信息與單個殘基內部的原子尺度結構信息的語義差異，研究只計算殘基內的 PDR，這也可以使 ESM-AA 學習到不同殘基內的各種結構知識。
在多尺度位置編碼 (Multi-scale Position Encoding) 方面，研究人員設計了一個多尺度位置編碼 E 來對代碼切換序列中的位置關系進行編碼。E 包含一個殘基尺度的位置編碼 ER 和一個原子尺度的位置編碼 EA。
對于 ER， 研究人員擴展了現有的編碼方法，使其能夠編碼從殘基到原子的關系，同時在處理純殘基序列時保持與原始編碼的一致性。對于 EA， 為了捕捉原子之間的關系，該研究直接使用空間距離矩陣 (spatial distance matrix) 對其三維位置進行編碼。
值得一提的是，多尺度編碼方法可以確保，預訓練不會受到模糊位置關系的影響，從而使 ESM-AA 在兩個尺度上都能有效地發揮作用。
在將多尺度 PE 集成到 Transformer 時，該研究首先用殘差尺度位置編碼 ER 替換了 Transformer 中的正弦編碼，將原子尺度的位置編碼 EA，視為自注意力層的偏置值 (bias term)。
研究結果：融合分子知識，優化蛋白質理解
為了驗證多尺度統一預訓練模型的有效性，該研究在各種涉及蛋白質和小分子的任務中評估了 ESM-AA 的表現。
如上表所示，在酶-底物親和力回歸任務、酶-底物對分類任務和藥物-靶點親和力回歸任務的性能比較中，大多數指標上，ESM-AA 都優于其他模型并達到了最先進的結果。 此外，微調策略（如 ProSmith 和 XGBoost）建立在 ESM-AA 上時，性能始終優于結合兩個獨立的分子預訓練模型與蛋白預訓練模型的版本（如表 1 和表 2 最后四行所示）。
值得注意的是，ESM-AA 甚至可以打敗使用了具有更大參數規模的預訓練模型的方法（如表 2 中第 5 行、第 7 行與最后一行的對比）。
為了驗證多尺度位置編碼的有效性，該研究在兩種情況下進行了消融實驗：一種是不使用原子尺度位置編碼 (ASPE) 的情況；另一種是不使用遺傳尺度位置編碼 (RSPE) 的情況下進行。
在刪除分子或蛋白質數據時，模型性能出現顯著下降。有趣的是，刪除蛋白質數據導致的性能下降比刪除分子數據更明顯。這表明，當模型沒有經過蛋白質數據訓練時，會迅速丟失與蛋白質相關的知識，從而導致整體性能明顯下降。然而，即使沒有分子數據，模型仍然可以通過解壓縮操作獲得原子級別的信息。
由于 ESM-AA 基于現有的 PLMs 開發，該研究希望確定其是否仍然保留了對蛋白質的全面理解，從而通過使用二級結構預測 (secondary structure prediction) 和無監督接觸預測 (unsupervised contact prediction) 任務，測試蛋白質預訓練模型在蛋白質結構理解方面的能力。
結果表明，雖然 ESM-AA 在此類研究中可能無法實現最佳性能，但其在二級結構預測和接觸預測方面與 ESM-2 的表現相似。
在分子基準測試中，ESM-AA 在大多數任務中與 Uni-Mol 表現相當， 并在許多情況下優于幾個特定于分子的模型，這表明其已成為一種處理分子任務的強大方法。
為了更直觀地說明 ESM-AA 獲得了更高質量的蛋白質和小分子表征，該研究在酶-底物對分類和藥物靶點親和力回歸任務中，對 ESM-AA 和 ESM-2+Uni-Mol 提取的表征進行了可視化比較。結果顯示，ESM-AA 模型能夠創建包含蛋白質和分子數據的更具內聚性的語義表征，這使得 ESM-AA 優于兩個單獨的預訓練模型。
蛋白質語言模型，大語言模型的下一段征程
大約從 1970 年代起，就有越來越多的科學家認為「二十一世紀是生物學的世紀。」去年 7 月，福布斯曾在一篇長文中暢想，LLM 令人們處于生物學領域新一輪變革的風口浪尖。生物學原來是一個可破譯、可編程，在某些方面甚至是數字化的系統，LLM 憑借其對自然語言的驚人駕馭能力，為破解生物語言提供了潛在可能， 這也讓蛋白質語言模型成為了這個時代最受關注的領域之一。
蛋白質語言模型代表了 AI 技術在生物學中的前沿應用。其通過學習蛋白質序列的模式和結構，能夠預測蛋白質的功能和形態，對于新藥開發、疾病治療和基礎生物學研究具有重大意義。
此前，蛋白質語言模型如 ESM-2 和 ESMFold 已經展現出與 AlphaFold 相媲美的準確性，并且具備更快的處理速度和對「孤兒蛋白質」更準確的預測能力。這不僅加速了蛋白質結構的預測，也為蛋白質工程提供了新的工具，使得研究人員能夠設計出具有特定功能的全新蛋白質序列。
此外，蛋白質語言模型的發展受益于所謂的「縮放法則」，即模型的性能隨著模型規模、數據集大小和計算量的增加而顯著提高。 這意味著，隨著模型參數的增加和訓練數據的積累，蛋白質語言模型的能力將得到質的飛躍。
近兩年，蛋白質語言模型在企業界也進入了快速發展時期。2023 年 7 月，百圖生科與清華大學聯合提出了一種名為 xTrimo Protein General Language Model (xTrimoPGLM) 的模型，參數量高達千億 (100B)，在多種蛋白質理解任務（15 項任務中的 13 項任務）中顯著優于其他先進基線模型。在生成任務上，xTrimoPGLM 能夠生成與自然蛋白質結構類似的新蛋白質序列。
2024 年 6 月，AI 蛋白質企業途深智合宣布，將其研發的國內首個自然語言蛋白質大模型 TourSynbio? 面向所有科研人員和開發者開源。 該模型以對話的方式實現了對蛋白質文獻的理解，包括蛋白質性質、功能預測和蛋白質設計等功能，在對比蛋白質評測數據集的測評指標上，超過 GPT4，成為行業第一。
此外，以 ESM-AA 為代表技術研究的突破，或許也意味著技術的發展即將度過「萊特兄弟時刻」，迎來飛躍。同時，蛋白質語言模型的應用也將不僅限于醫療和生物制藥領域，還可能擴展到農業、工業、材料科學和環境修復等多個領域，推動這些領域的技術革新，為人類帶來前所未有的變革。
