ICML 2024｜用於統一分子建模的多尺度蛋白質語言模型ESM-AA

实验盒發表於2024-06-06
模型

多模態蛋白質語言模型是一個新興的方向。 《Multi-scale Protein Language Model for Unified Molecular Modeling》是一篇發表在ICML 2024的文獻，作者是 Kangjie Zheng、Siyu Long、Tianyu Lu、Junwei Yang、Xinyu Dai、Ming Zhang、Zaiqing Nie、Wei-Ying Ma 和 Hao Zhou。文章的蛋白質語言模型涉及到氨基酸序列和分子結構的多模態資訊，提供了一個整合的思路。

摘要

蛋白質語言模型在蛋白質工程領域展現出巨大潛力。然而，現有的蛋白質語言模型主要在殘基（氨基酸）尺度上運作，這限制了它們提供原子尺度資訊的能力。這種限制阻礙了我們充分利用蛋白質語言模型在涉及蛋白質和小型分子的應用程式中的潛力。因此，作者提出了一種新方法ESM-AA（ESM All-Atom），它透過預訓練多尺度程式碼切換蛋白質序列，並使用多尺度位置編碼來捕捉殘基和原子之間的關係，從而實現了原子尺度和殘基尺度統一的分子建模。實驗結果表明，ESM-AA在蛋白質分子任務中超越了先前的方法，展示了蛋白質語言模型的充分利用。進一步的調查揭示了透過統一分子建模，ESM-AA不僅獲得了分子知識，還保留了對蛋白質的理解。

引言

蛋白質語言模型在蛋白質工程中展現出巨大潛力，它們在預訓練大規模蛋白質序列時能夠捕獲生化和共同進化知識。這在多個領域取得了顯著成就，包括蛋白質結構預測、蛋白質適應度預測和蛋白質設計等。然而，現有的蛋白質語言模型主要在殘基尺度上運作，不提供原子尺度的資訊。為了充分利用蛋白質語言模型在涉及大分子（蛋白質）和小型分子的應用程式中的潛力，需要將外部小型分子模型納入考慮。作者提出了ESM-AA，它透過預訓練多尺度程式碼切換蛋白質序列，並使用多尺度位置編碼來描述殘基和原子之間的關係，實現了多尺度統一分子建模。

方法

文獻提出了提出一個多尺度預訓練模型ESM-AA，它透過以下兩個主要步驟實現多尺度統一分子建模：1. 在多尺度程式碼切換蛋白質序列上進行預訓練；2. 使用多尺度位置編碼描述殘基和原子之間的關係。

image

多尺度預訓練過程概述

模型架構的關鍵特點和元件如下：

  1. 多尺度預訓練（Multi-scale Pre-training）:

    • 程式碼切換蛋白序列（Code-Switch Protein Sequence）: 模型透過隨機“解壓縮”（unzipping）部分殘基（氨基酸）到它們的組成原子，併為每個解壓縮的原子分配座標，從而在蛋白質序列中引入原子級別的資訊。
    • 多尺度位置編碼（Multi-scale Position Encoding, MSPE）: 為了在程式碼切換序列中描述殘基和原子之間的複雜位置關係，ESM-AA設計了一種多尺度位置編碼，包括殘基尺度位置編碼（Residue Scale Position Encoding, RSPE）和原子尺度位置編碼（Atom Scale Position Encoding, ASPE）。

  2. 殘基尺度位置編碼（Residue Scale Position Encoding, RSPE）:

    • 基於現有的編碼方法，如Rotary Position Embedding (RoPE)，用於描述殘基之間的相對位置關係。
    • 對於同一殘基中的原子，重用該殘基的位置編碼，避免引入模糊的位置資訊。

  3. 原子尺度位置編碼（Atom Scale Position Encoding, ASPE）:

    • 使用空間距離矩陣和高斯核（Gaussian Kernel）直接對原子的3D位置進行編碼，以描述原子之間的關係。

  4. 變換器編碼器（Transformer Encoder）:

    • ESM-AA採用了標準的Transformer架構，其中位置編碼被整合到自注意力層中。
    • 自注意力計算時，原子尺度位置編碼被視作自注意力層的偏置項（bias term）。

  5. 預訓練任務（Pre-training Objectives）:

    • 掩碼語言建模（Masked Language Modeling, MLM）: 模型需要預測被掩蓋（masked）的殘基或原子。
    • 成對距離恢復（Pair-wise Distance Recovery, PDR）: 模型需要從被噪聲汙染的原子座標中恢復準確的原子間歐幾里得距離。

  6. 多尺度掩碼語言建模（Multi-scale MLM）:

    • 模型同時對殘基尺度和原子尺度的標記進行掩碼和預測。

  7. 成對距離恢復（Pair-wise Distance Recovery, PDR）:

    • 模型需要從被噪聲汙染的原子座標中恢復準確的原子間距離。

  8. 模型引數化（Parameterization）:

    • ESM-AA使用12層堆疊的Transformer層，每層有20個注意力頭。
    • 模型維度和前饋維度分別為480和1920。

  9. 輸入處理:

    • 模型可以接受蛋白質或分子作為輸入，但在預訓練中，輸入是未配對的蛋白質或分子資料。

  10. 輸出:

    • 模型輸出可以用於多種下游任務，如蛋白質結構預測、蛋白質適應度預測、蛋白質設計等。

image

包括多尺度掩蔽語言建模和成對距離恢復的多尺度預訓練框架

image

多尺度位置編碼框架

ESM-AA模型架構的創新之處在於它能夠同時處理蛋白質的殘基尺度和原子尺度資訊，並透過統一的建模方法提高了蛋白質-分子任務的效能。這種多尺度建模策略使得ESM-AA能夠更全面地理解和預測蛋白質及小分子的結構和功能。

實驗結果

作者在蛋白質和小型分子的混合資料上預訓練ESM-AA，並在多種基準測試上對其進行微調，以驗證其效能。主要結果包括：

  1. 效能提升: ESM-AA在蛋白質分子任務中的表現超越了以往的方法。這表明透過統一的分子建模方法，可以充分利用蛋白質語言模型的潛力。

  2. 多尺度建模: ESM-AA成功實現了在原子尺度和殘基尺度上的統一建模。透過預訓練多尺度程式碼切換蛋白質序列，並使用多尺度位置編碼，模型能夠同時捕獲殘基間和原子間的關係。

  3. 實驗驗證: 透過一系列實驗，作者驗證了ESM-AA在多種基準測試上的有效性。這些基準測試包括蛋白質結構預測、蛋白質適應度預測、蛋白質設計等任務。

  4. 保留蛋白質理解: 儘管ESM-AA被設計為多尺度模型，它仍然保留了對蛋白質的深入理解。這透過在蛋白質特定任務（如二級結構預測和接觸預測）上與ESM-2模型相似或更好的效能來證明。

  5. 分子任務的表現: 當將ESM-AA應用於標準分子基準測試時，它也超越了幾個特定於分子的模型，這突出了統一分子建模的潛力。

  6. 視覺化分析: 透過對ESM-AA和ESM-2+Uni-Mol學習到的表示進行視覺化比較，結果表明ESM-AA模型能夠建立更一致的語義表示，包括蛋白質和分子資料。

  7. 在虛擬篩選基準測試中的效能: ESM-AA在虛擬篩選基準測試中表現出強大的效能，即使在零樣本設定下，也接近最先進的方法DrugCLIP。

  8. 蛋白質功能註釋任務的效能: ESM-AA在蛋白質功能註釋任務上的表現與結構蛋白質表示模型相當或更好，即使沒有結構資訊輸入。

  9. 消融研究: 文獻中的消融研究表明，多尺度位置編碼、預訓練目標和訓練資料等各個組成部分對ESM-AA方法的有效性至關重要。

這些結果表明，ESM-AA模型不僅在理論上是創新的，而且在實際應用中也是有效的，能夠處理蛋白質和分子的複雜互動，並在多種生物學和化學任務中表現出色。

總結

蛋白質語言模型已廣泛應用於多個領域，包括蛋白質結構預測、蛋白質適應度預測和蛋白質設計。統一分子建模擴充套件了蛋白質語言模型在殘基和原子尺度上有效運作的能力，從而增強了它們在這些任務中的適用性。此外，統一分子建模也為蛋白質-小分子相互作用領域的研究開闢了新的途徑。基於此，作者提出了ESM-AA，這是一個多尺度蛋白質語言模型，透過預訓練多尺度程式碼切換蛋白質序列和使用多尺度位置編碼來描述殘基和原子之間的關係，實現了多尺度統一分子建模。實驗結果表明，ESM-AA在蛋白質-分子任務中超越了先前的方法，並且有效地將分子知識整合到蛋白質語言模型中，同時沒有犧牲對蛋白質的理解。

相關連結

[1] 文獻：https://arxiv.org/abs/2403.12995v2

[2] 審稿意見：https://openreview.net/pdf?id=MBIGXMT0qC

相關文章