編輯 | KX
AI 已廣泛用於藥物發現和材料設計中的分子篩選。當前的自監督分子預訓練方法往往忽略了分子的基本化學特性和物理原理。
為了解決這個問題，來自清華大學的研究人員提出了一種稱為分數去噪（Fractional denoising，Frad）的分子預訓練框架，以學習有效的分子表示。透過這種方式，噪聲變得可自定義，允許納入化學先驗，從而大大改善分子分佈建模。
實驗表明，Frad 始終優於現有方法，在力預測、量子化學特性和結合親和力任務中，取得最先進的結果。改進的噪聲設計提高了力準確性和取樣覆蓋範圍，這有助於建立物理一致的分子表示，最終實現卓越的預測效能。
相關研究以「Pre-training with fractional denoising to enhance molecular property prediction」為題，於 9 月 18 日釋出在《Nature Machine Intelligence》上。
論文連結：https://www.nature.com/articles/s42256-024-00900-z
AI 用於分子性質預測的挑戰
分子性質預測是藥物發現和材料設計等各個領域的一項關鍵任務。傳統方法，包括第一性原理計算和溼實驗室實驗，需要巨大的成本，從而阻礙了具有所需特性的分子的高通量篩選。科學家認為深度學習方法可以降低成本和大幅加速篩選過程。
深度學習分子性質預測方法面臨的主要困難是標記分子資料的稀缺性。為了緩解這一困難，已經提出了各種自監督分子預訓練方法來利用未標記分子資料中的內在資訊。
不幸的是，直接借用計算機視覺和自然語言處理中流行的預訓練任務可能不適合分子，因為它們忽略了分子的基本物理原理。
相比之下，在預訓練中應用去噪可以解釋為等效的力學習，但有限的噪聲分佈會給分子分佈帶來偏差。這裡的挑戰在於有效地模擬全面的分子分佈，同時保留力學習的基本物理解釋。
Frad：分數去噪的分子預訓練框架
鑑於模擬真實分子分佈的難度，清華團隊選擇透過將有關分子分佈的化學先驗引入噪聲設計來更全面地表徵分佈，由於噪聲分佈的限制，這在以前的方法中是被禁止的。
因此，研究人員提出了一種稱為分數去噪（Frad）的分子預訓練框架，證明其具有力學習解釋。
具體而言，給定一個平衡分子構象，新增化學感知噪聲（CAN）和座標高斯噪聲（CGN）的混合噪聲並獲得噪聲構象。訓練模型以從噪聲構象中預測 CGN。「分數」（fractional）是指恢復引入的全部噪聲的一小部分。特別是，CAN 是可定製的，使 Frad 能夠結合化學先驗來最佳化分子分佈建模。
受描述分子構象變化的化學先驗的啟發，研究人員提出了兩個版本的 CAN。具體來說，旋轉噪聲（RN）被用來捕捉單鍵的旋轉，而振動和旋轉噪聲（VRN）被用來反映各向異性的振動。
Frad 框架的整個流程如下圖所示。給定一個平衡分子構象 xeq，新增 CAN 和 CGN 的混合，其中平衡構象是指分子勢能面區域性最小值處的結構。
然後，訓練模型從噪聲構象（即 Frad）中預測 CGN，因為它恢復了一部分引入的噪聲。具體來說，利用等變圖神經網路 (GNN) 從噪聲構象中提取特徵，然後噪聲預測頭根據這些特徵預測 CGN。
研究使用 TorchMD-NET 作為骨幹模型，從 3D 分子輸入中獲得等變特徵。具體來說，理論分析表明，無論 CAN 分佈如何，該任務都具有力學習解釋，而 CAN 分佈會影響力目標和取樣分佈。因此，精心設計 CAN 以與真實的分子構象分佈保持一致，與現有的去噪方法相比，可以實現更精確的力目標和更廣泛的取樣分佈。
在微調過程中，研究人員從預訓練權重初始化等變 GNN，隨後為每個單獨的下游任務量身定製不同的屬性預測頭。預訓練的 GNN 權重以及預測頭中的引數在下游標籤的監督下繼續更新。
圖示：Frad 框架概述。（來源：論文）
Frad 提升了性質預測效能
為了評估 Frad 在預測分子性質方面的有效性，進行了一系列具有挑戰性的下游任務，包括原子級力預測、分子級量子化學性質預測和蛋白質-配體複合物級結合親和力預測。
Frad 模型與已建立的基線進行了系統比較，包括預訓練方法以及未經預訓練的性質預測模型。
原子級力預測
首先，評估了 Frad 在各種分子動力學資料集（即 MD17、ISO17 和 MD22）中預測力的表現。
Frad 在八種分子中的七種上的表現都優於基線。特別是與 Frad 方法具有相同主幹的去噪預訓練方法，即 3D-EMGP 和 Coord 相比，Frad 的優越性顯而易見。這凸顯了在去噪中引入化學先驗以準確模擬分子分佈的必要性。在訓練集的大小有限的情況下，Frad 預訓練始終能大幅度增強主幹模型 TorchMD-NET，表明 Frad 有效地學習了分子的內在原理並具有出色的泛化能力。
對於大訓練集設定中的苯，在 Frad 的微調過程中出現了過擬合現象，而這在其他分子中並沒有觀察到。這可能是由於苯的結構相對固定，導致低維特徵容易出現過擬合。
ISO17 和 MD22 上的結果如表 1 所示。骨幹模型 TorchMD-NET 的效能透過 Frad 預訓練不斷提高。
分子級量子化學性質預測
為了驗證 Frad 是否能在不同性質上取得有競爭力的結果，研究人員在 QM9 資料集中的 12 個任務上評估了 Frad(RN) 和 Frad(VRN)。表 2 總結了 Frad 的結果以及預訓練和監督基線。
總體而言，Frad 超越了監督和預訓練方法，並在 12 個目標中的 9 個目標上取得了最先進的結果。此外，在 11 個目標上基於骨幹模型 TorchMD-NET 取得了顯著的改進，表明 Frad 預訓練學到的知識對大多數下游任務都有幫助。
此外，在表格下半部分明顯優於去噪預訓練方法。特別是，Frad 在所有 12 個任務中都達到或超過了使用相同骨幹 TorchMD-NET 的 Coord 方法的結果，表明引入的分佈化學先驗有助於多類性質。還注意到 VRN 和 VN 的效能相當。以此推測，對於 QM9 資料集中對輸入構象不太敏感的任務（與力預測任務相比），各向異性振動噪聲的建模可能並不重要。
複合物結合親和力預測
蛋白質-配體結合親和力 (LBA) 預測旨在預測蛋白質和配體之間的相互作用強度。研究人員在 ATOM3D 資料集中評估了 Frad 在 LBA 任務上的表現。
在 ATOM3D 預處理之後，Frad 從蛋白質中提取結合口袋，並利用口袋-配體複合物結構作為輸入來預測結合親和力值。結果顯示在表 3 中。
Frad 擊敗了針對蛋白質表示量身定製的基於序列和基於結構的基線，表明 Frad 是一種適用於小分子和蛋白質的通用表示學習方法。預訓練基線 SE(3)-DDM 不評估訓練集和測試集之間涉及 60% 蛋白質序列同一性的分割設定，是一種等變 Coord 方法，可對成對原子距離上的高斯噪聲進行去噪。同樣，Frad 超越了傳統的去噪方法，再次表明了引入化學先驗的重要性。
未來工作
未來研究的幾個探索方向如下：
- 首先，增加預訓練資料量有可能顯著提高整體效能。由於獲得精確的分子構象成本高昂，目前使用的預訓練資料集仍然比二維和一維分子資料集小得多。將來會有更多 3D 分子資料可用。
- 其次，目前的重點是使用 3D 輸入進行屬性預測。透過與其他預訓練方法相結合，可以生成一個能夠跨資料模態處理分子任務的模型。
- 最後，如何為典型的分子類別設計 CAN 值得研究，例如核酸、蛋白質和材料，以便 Frad 可以有效地應用於更廣泛的領域，並加快藥物和材料的發現。
透過這些改進，Frad 有可能成為適用於各種分子任務的強大分子基礎模型。從而催化藥物發現和材料科學等領域的突破。