來自清華大學AIR、北京大學、南京大學的研究團隊提出了 ESM-AA 模型。該模型在蛋白質語言建模領域取得了重要進展，提供了一套整合多尺度資訊的統一建模方案。
它是首個能同時處理氨基酸資訊和原子資訊的蛋白質預訓練語言模型。模型的出色效能展示了多尺度統一建模在克服現有侷限和解鎖新能力方面的巨大潛力。
作為基座模型，ESM-AA 獲得了多位學者的關注與廣泛討論（截圖見下方），被認為有潛力基於 ESM-AA 開發出可與 AlphaFold3、RoseTTAFold All-Atom 相競爭的模型，為研究不同生物結構間的相互作用開闢了新的道路。當前論文已被 ICML 2024 錄⽤。
研究背景
蛋白質是各種生命活動的關鍵執行者。深入理解蛋白質及其與其他生物結構的相互作用是生物科學中的核心議題，這對靶向藥物篩選、酶工程等領域具有顯著的實際意義。
因此，如何更好地理解與建模蛋白質也成為了目前 AI4Science 領域的一個研究熱點。
近日來，包括 Deepmind、華盛頓大學 Baker 組在內的各大前沿研究機構也針對蛋白質全原子建模問題展開了深入研究，提出了包括 AlphaFold 3、RoseTTAFold All-Atom 等針對蛋白質以及其他生命活動相關分子的全原子尺度建模模型，可以在很高的精度下實現對蛋白質結構、分子結構以及受體-配體結構等全原子尺度的精確預測。
雖然這些模型對於全原子尺度的結構建模取得了重大進展，目前主流的蛋白語言模型仍然無法實現全原子尺度的蛋白質理解與表示學習。
多尺度，下一代蛋白質模型的「必經之路」
以 ESM-2為代表的蛋白質表示學習模型，它們以氨基酸作為構建模型的唯一尺度，這對於專注於處理蛋白質的情境而言是一種合理的方法。
然而，要全面理解蛋白質的本質，關鍵在於闡述它們與其他生物結構（如小分子、DNA、RNA 等）之間的相互作用。
面對這種需求，需要描述不同結構間複雜的相互作用，單一尺度的建模策略難以提供有效的全面覆蓋。
為了克服這一缺陷，蛋白質模型正在經歷一場向多尺度模型轉變的深刻革新。例如，5月初發表在《Science》雜誌上的 RoseTTAFold All-Atom 模型，作為 RoseTTAFold 的後續產品，引入了多尺度概念。
這一模型不僅僅侷限於蛋白質結構預測，還擴充到了蛋白質與分子/核酸的對接、蛋白翻譯後修飾等更廣泛的研究領域。
同時，DeepMind 最新發布的 AlphaFold3 也採用了多尺度建模策略，支援預測多種蛋白複合物的結構，其表現令人矚目，無疑將對人工智慧和生物學領域產生重大影響。
ESM All-Atom，多尺度的蛋白質語言模型基座
RoseTTAFold All-Atom 和 AlphaFold3 對多尺度概念的成功應用啟發了一個重要思考，即：作為蛋白質基座模型的蛋白語言模型應如何採納多尺度技術。基於此，該團隊提出了多尺度蛋白質語言模型 ESM All-Atom（ESM-AA）。
簡要來說，ESM-AA 透過將部分氨基酸「展開」（Unzip）為對應的原子組成引入了多尺度概念。隨後，透過混合蛋白資料與分子資料進行預訓練，這使得模型具備了同時處理不同尺度生物結構的能力。
此外，為了幫助模型更好地學習優質的原子尺度資訊，ESM-AA 還會利用原子尺度的分子結構資料進行訓練。而且透過引入圖 2 所示的多尺度位置編碼機制，ESM-AA 模型可以很好地對不同尺度的資訊進行區分，確保模型能夠精確理解殘基層面與原子層面的位置與結構資訊。
多尺度預訓練目標
為了幫助模型學習多尺度資訊，該團隊為 ESM-AA 模型設計了多種預訓練目標。ESM-AA 的多尺度預訓練目標包括掩碼語言建模（MLM）和成對距離恢復（PDR）。如圖 3(a) 所示，MLM 透過遮蓋氨基酸和原子，要求模型根據周圍的上下文進行預測，這一訓練任務可以在氨基酸和原子兩個尺度上進行。而 PDR 則要求模型準確預測不同原子之間的歐幾里得距離，以訓練模型理解原子級的結構資訊（如圖 3(b) 所示）。
實驗驗證
效能評估
ESM-AA 模型在多種蛋白-小分子基準任務上進行微調和評估，包括酶-底物親和力迴歸任務（結果展示於圖 4）、酶-底物對分類任務（結果展示於圖 4）和藥物-靶標親和力迴歸任務（結果展示於圖 5）。
結果顯示，ESM-AA 在這些任務中優於之前的模型，表明其在氨基酸和原子尺度上充分發揮了蛋白質預訓練語言模型的潛力。
圖 5：藥物-靶標親和力迴歸任務效能比較
此外，ESM-AA 模型也在蛋白質接觸預測、蛋白功能分類以及分子性質預測等任務上測試了效能。
結果顯示，在處理僅涉及蛋白質的任務時，ESM-AA 的表現與 ESM-2 相當；在分子任務上，ESM-AA 模型的效能優於大多數基準模型，與 Uni-Mol 的表現相近。
這表明，ESM-AA 在獲取強大分子知識的過程中並未犧牲對蛋白質的理解能力，也進一步說明了 ESM-AA 模型成功複用了 ESM-2 模型的知識，而無需被從頭開始重新開發，顯著降低了模型訓練成本。
視覺化分析
為了進一步分析 ESM-AA 在蛋白-小分子基準任務上表現優異的原因，該論文展示了 ESM-AA 模型以及 ESM-2+Uni-Mol 模型組合在該任務中抽取的樣本表徵分佈的視覺化結果。
如圖 6 所示，ESM-AA 模型所學習的蛋白和小分子表示之間更為緊湊，這預示著兩者處於同一表示空間，這是 ESM-AA 模型優於 ESM-2+Uni-Mol 模型的原因，進一步說明了多尺度統一分子建模的優勢。
結語
清華 AIR 團隊開發的 ESM-AA 是首個融合氨基酸與原子資訊處理的蛋白質預訓練語言模型。模型透過整合多尺度資訊，展現出穩健且卓越的效能，為解決生物結構間相互作用難題提供了新途徑。
ESM-AA不僅促進了對蛋白質更深層次的理解，還在多項生物分子任務中表現出色，證明了其在保持蛋白質理解能力的同時可以有效融合分子層面知識，降低了模型訓練的成本，為 AI 輔助的生物科學研究開闢了新方向。
論文標題：ESM All-Atom: Multi-Scale Protein Language Model for Unified Molecular Modeling