清華AIR等提出ESM-AA，首個從胺基酸到原子尺度的蛋白質語言模型

來自清華大學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，多尺度的蛋白質語言模型基座

圖1：ESM-AA 模型的整體架構

RoseTTAFold All-Atom 和AlphaFold3 對多尺度概念的成功應用啟發了一個尺度概念的成功應用啟發了一個尺度概念的成功應用啟發了一個尺度概念的成功應用啟發重要思考，即：作為蛋白質基座模型的蛋白質語言模型應如何採納多尺度技術。基於此，團隊提出了多尺度蛋白質語言模型 ESM All-Atom（ESM-AA）。

簡要來說，ESM-AA 透過將部分胺基酸「展開」（Unzip）為對應的原子組成引入了多尺度概念。隨後，透過混合蛋白質資料與分子資料進行預訓練，這使得模型具備了同時處理不同尺度生物結構的能力。

此外，為了幫助模型更好地學習優質的原子尺度信息，ESM-AA 還會利用原子尺度的分子結構數據進行訓練。而且透過引入圖 2 所示的多尺度位置編碼機制，ESM-AA 模型可以很好地對不同尺度的資訊進行區分，確保模型能夠精確理解殘基層面與原子層面的位置與結構資訊。

圖 2：多尺度位置編碼

多尺度預訓練目標

為了幫助模型學習多尺度訊息，該團隊為 ESM-AA 模型設計了多種預訓練目標。 ESM-AA 的多尺度預訓練目標包括遮罩語言建模（MLM）和成對距離恢復（PDR）。如圖 3(a) 所示，MLM 透過遮蓋胺基酸和原子，要求模型根據周圍的上下文進行預測，這項訓練任務可以在胺基酸和原子兩個尺度上進行。而 PDR 則要求模型準確預測不同原子之間的歐幾里德距離，以訓練模型理解原子級的結構資訊（如圖 3(b) 所示）。

圖3：多尺度預訓練任務

實驗驗證

性能評估

ESM-AA 模型在多種蛋白-小分子基準上進行微調和小分子基準評估物親和力回歸任務（結果顯示於圖4）、酵素-底物對分類任務（結果顯示於圖4）和藥物-標靶親和力回歸任務（結果顯示於圖5）。

結果顯示，ESM-AA 在這些任務中優於先前的模型，顯示其在胺基酸和原子尺度上充分發揮了蛋白質預訓練語言模型的潛力。

圖4：酶-底物親和力回歸任務以及酶-底物對分類任務性能比較

圖5：藥物-靶標親和力回歸任務性能比較在蛋白質接觸預測、蛋白質功能分類以及分子性質預測等任務上測試了表現。

結果顯示，在處理僅涉及蛋白質的任務時，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 模型的原因，進一步說明了多尺度統一分子建模的優勢。

圖 6：蛋白/分子表示的可視化分析

結語

清華 AIR 團隊開發的 ESM-AA 是首個融合氨基酸與原子訊息處理的蛋白質預訓練模型。模型透過整合多尺度訊息，展現出穩健且卓越的性能，為解決生物結構間相互作用難題提供了新途徑。

ESM-AA不僅促進了對蛋白質更深層次的理解，還在多項生物分子任務中表現出色，證明了其在保持蛋白質理解能力的同時可以有效融合分子層面知識，降低了模型訓練的成本，為AI 輔助的生物科學研究開啟了新方向。

論文標題：ESM All-Atom: 

Multi-Scale Protein Language Model for Unified Molecular Modeling

Github開源位址：

https://github.com/zhengkangange/ESM-AA ://arxiv.org/abs/2403.12995

以上是清華AIR等提出ESM-AA，首個從胺基酸到原子尺度的蛋白質語言模型的詳細內容。更多資訊請關注PHP中文網其他相關文章！