作為細胞内無數生化反應的驅動力，蛋白質在細胞微觀世界中扮演着建築師和工程師的角色，不僅催化着生命活動，更是構築、維系生物體形态與功能的基礎構件。正是蛋白質之間的互動、協同作用，支撐起了生命的宏偉藍圖。
然而，蛋白質的結構複雜多變，傳統的實驗方法在解析蛋白質結構時既耗時又費力——蛋白質語言模型 (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 的表現。
表1：在酶-底物親和力回歸任務 (ESAR) 和酶-底物對分類任務 (ESPC) 上的性能比較
表2：藥物-靶點親和力回歸任務的性能比較
如上表所示，在酶-底物親和力回歸任務、酶-底物對分類任務和藥物-靶點親和力回歸任務的性能比較中，大多數名額上，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-2+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 為代表技術研究的突破，或許也意味着技術的發展即将度過「萊特兄弟時刻」，迎來飛躍。同時，蛋白質語言模型的應用也将不僅限于醫療和生物制藥領域，還可能擴充到農業、工業、材料科學和環境修複等多個領域，推動這些領域的技術革新，為人類帶來前所未有的變革。