論文標題：Automated 3D Pre-Training for Molecular Property Prediction

論文鏈接：https://arxiv.org/abs/2306.07812

代碼鏈接：https://github.com/LARS-research/3D-PGT

該論文已經被KDD 2023接收。

1.  簡介

分子性質預測是藥物發現和材料科學等領域中的一個重要問題。考慮到分子的3D結構信息與其性質緊密相關，近年來，一個主流的研究熱點是將分子的3D結構信息與各種以圖學習為代表的機器學習方法結合，以提高分子性質的預測性能。然而，由于高昂的計算成本，在大規模數據實時計算分子的3D結構幾乎是不可行的[1]。

圖1：cc(=o)oc1=cc=cc=c1c(=o)o的2D分子圖（左）和3D構象圖（右）。

面對這一挑戰，我們提出了一個以預訓練范式為基礎的大模型，專門用于分子性質預測。預訓練范式隸屬于 AIGC 的一個重要技術分支[2]，其能處理大規模數據集并提取深層次特征，充分利用了大模型的優勢。本文在分子領域結合已有的3D分子構象進行預訓練，并在實際下游任務中只基于分子的2D結構信息進行微調并進行性質預測，從而在保證高效的前提下進一步提升在下游任務上的表現。

在這項工作中，我們提出了一種自動化3D預訓練框架：3D-PGT。基于分子的化學鍵長、鍵角和二面角是對應于完整分子3D構象的三個基本幾何描述符這一事實，我們對應設計了三個生成式預訓練任務，使得模型通過預訓練能夠具備編碼3D幾何結構的能力。而為了自動化分配這三個預訓練任務的權重以融合成一個總體的預訓練目標函數，我們基于分子總能量設計了一個surrogate metric，從而自動搜索三個預訓練任務的權重分布。

圖2：分子性質預測問題的核心問題示意：如何兼具更快的推理速度和更小的推理誤差。

為了驗證所設計的預訓練框架的有效性，我們基于已測定 3D 結構的公開分子數據集進行了預訓練，并在 8 個量子化學性質預測任務和 12 個涉及藥理學、分子化學的下游任務上進行微調和實驗驗證。結果表明，3D-PGT能通過3D預訓練帶來明顯的性能增益，且優于其他預訓練 baseline。這再次證明了在分子性質預測領域，大模型處理復雜任務和大規模數據的優勢。同時，基于 3D-PGT 的方法還在催化劑仿真挑戰 Open Catalyst Challenge 2022 上取得第 3，在大規模分子性質預測挑戰 OGB-LSC@NeruIPS challenge 2022 中取得第 8 名的成績。

2.  背景介紹

分子性質作為描述分子特征和行為的基本屬性，在藥理學、材料化學等研究領域和具體應用都有著重要意義。然而分子性質的測定往往需要借助成本高昂的計算方式，例如以密度泛函理論（Density Functional Theory, DFT）為代表的計算化學方法。近年來，隨著大規模量子化學計算和高通量實驗的技術進步，一個新興的工業界和學術界熱門的研究方向，是利用具有適當歸納偏置的機器學習方法并結合已測定性質的大量分子數據，實現高效的分子性質預測，并應用于一系列實際下游應用中，例如大規模藥物分子篩選，合成材料篩選，催化劑設計等。

圖3：分子性質預測效率對比示意圖。其中基于量子化學計算的DFT計算特定分子性質需要數個小時，而利用機器學習方法只需要遠小于1秒的時間。該圖出自OGB@NeurIPS 2022 Challenge，預測分子的HOMO-LUMO能隙。

在過去的幾年中，一種主流的方法是將分子建模為 2D 圖結構，其中以原子作為節點，而邊作為化學鍵，將整個分子性質預測建模為圖級（Graph-level）預測任務，并應用 GNN（Graph Neural Network）通過擬合 DFT 的計算方式預測分子的性質。但目前，這種方法只有較高的預測效率，預測性能距離實際應用還有較大的差距。

圖4：現有分子性質預測方法的一般流程。基于DFT的量子化學計算方法首先對分子的三維結構進行優化以獲得低能構象，然后基于構象計算特定的分子性質。2D模型近似DFT的整個過程，通過輸入的2D分子圖直接預測分子性質；3D模型在2D分子圖的基礎上還需要進一步輸入3D分子構象，以獲得更好的預測性能。

由于分子的 3D 結構反映了原子和官能團之間的相互作用和相對位置，因此 3D 結構對分子性質的理解至關重要。考慮到這一點，一系列方法針對分子的 3D 結構信息設計 3D 模型，從而獲得更好的預測效果。然而，3D 結構的獲取需要依賴 DFT 等量子化學計算手段，這種昂貴的計算成本導致分子的 3D 結構在許多實際的下游任務中往往是不可獲取的。

結合上述兩種主流路線各自的優缺點，一個最近比較新穎的角度是將包含已測定 3D 結構信息的數據集作為預訓練數據集，基于 3D 結構設計預訓練任務，從而讓模型理解 DFT 由 2D 分子圖計算優化 3D 結構的過程，并將該先驗信息遷移到下游 2D 分子性質預測任務上，從而兼具效率和性能。目前，主流的同期工作例如 GraphMVP[3] 和3D Infomax[4] 都通過對齊分子 2D 視圖和 3D 視圖的圖級表示向量來設計預訓練目標函數，并通過在主流 benchmark 上的實驗證明了 3D 預訓練方案的有效性。

綜上，本文繼續沿著分子 3D-Pretraining 的 Pipeline，設計了三個基礎的生成式預訓練任務，并基于分子總能量設計目標函數來自動搜索各預訓練任務的權重，從而構建多預訓練任務的自動融合框架，并在廣泛的下游任務上獲得了顯著的預訓練收益。

3.  本文的方法

3.1  生成式預訓練任務的設計

DFT 優化分子 3D 結構的核心是在勢能面上尋找局部最小值[5]，分子能量出于勢能面上局部最小值時的 3D 結構也被稱為分子構象。其中，構象所包含的 3D 幾何信息可以被以下三個描述符完整地描述：兩個原子之間的化學鍵長、兩個化學鍵之間形成的夾角，以及三個化學鍵形成的二面角。

生成式預訓練任務使模型能夠理解分子從 2D 拓撲到 3D 幾何的基于 DFT 的優化過程。由于分子幾何是由電子的量子力學行為決定的，因此生成式預訓練任務可以通過學習 3D 構象的生成來間接學習量子化學性質的預測。基于此，本文設計了三個生成式預訓練任務來分別生成這鍵長、鍵角、二面角這三個描述符，具體如下所示：

圖5：分子結構和勢能面關系示意圖。勢能面將分子能量定義為多個坐標軸的函數，分子3D構象的優化就是在勢能面上尋找局部最小值，而3D構象中的幾何信息可以通過連接兩個原子的鍵長、三個相連原子的鍵角和三個連續鍵的二面角來描述

其中，h 代表對應節點的表示向量，f 為對應的預測網絡（這里的設置是MLP）。可以看到，這三個任務的核心思路都是利用 backbone 提取的節點表示來設計具體的回歸任務，通過將描述符中包含的 3D 結構信息當做解讀信號，從而讓 backbone 具有編碼 3D 結構信息的能力，并將這種能力作為先驗嵌入到模型并遷移到實際的下游任務中。

同時，考慮到當鄰居數為|N|時，分子的鍵角和二面角的計算復雜度呈 Ο(|N|^2 )和Ο(|N|^3 ) 增長，使得大規模的 3D 預訓練實現成本過高。本文基于 RGC（Runtime Geometry Calculation）重新設計了鍵角和二面角的目標函數，用每個原子的所涉及的鍵角和與每個化學鍵設計的二面角的和來代替所有鍵角和二面角的預測，從而將計算復雜度降低到線性級別。

3.2  自動化多預訓練任務融合框架

由于我們同時定義了多個預訓練任務，且各預訓練任務的損失函數對于模型參數的梯度優化方向并不一致，因此各自預訓練任務的權重分配是一個需要考慮的問題。如何自動化且最更好地分配多預訓練任務的權重，以預期在下游任務上獲得更好的性能增益，是設計多預訓練任務融合框架的核心。

由于鍵長、鍵角和二面角都是局部描述符，因此這里需要一個明確的 surrogate metric 來評估預訓練后的 backbone 對于分子 3D 結構整體的編碼能力而不是局部編碼能力。對此，本文考慮得到分子總能量和分子 3D 結構之間的對應關系，基于分子總能量設計了一個目標函數，并基于此設計了 bi-level 的油畫框架來搜索三個生成式預訓練任務各自的權重。最終，我們整體預訓練的 pipeline 如下圖所示。在預訓練階段，我們首先基于預訓練分子的低能構象設計了三個生成式預訓練任務，并基于分子低能構象對應的分子總能量設計了一個 surrogate metric 來搜索三個預訓練任務的各自權重；而在下游的微調階段，由于我們已經在預訓練階段引入了分子幾何先驗，因此通過微調即可在實際下游任務中獲得性能增益。

圖6：3D-PGT的預訓練框架示意圖

4.  實驗結果

在本文中，我們針對性的設計了一系列實驗來證明 3D 預訓練這一技術路線和本文設計的預訓練任務及自動融合框架的有效性，從而面臨當只有 2D 分子圖可以用于預測時，可以避免為每個分子生成 3D 構象的巨大計算成本，并能通過微調獲得明顯的性能收益。

本文主要在 3 個包含 3D 結構信息的數據集上進行預訓練，并主要在 12 個下游任務上進行微調和效果驗證，數據集的詳細統計數據如下所示：

圖7：數據集的統計細節

4.1 量子化學性質預測的表現

本文首先在流行的基準數據集QM9上評估了3D-PGT的量子化學領域的性質預測能力。我們首先從QM9包含的134k單一分子構象樣本中隨機選取50k個攜帶構象信息的分子樣本用于預訓練，并從剩余的樣本中選取50k個分子樣本并屏蔽其3D結構信息用于微調和評測。結果如下：

圖8：QM9數據集實驗對比結果，評價指標為MAE（平均絕對誤差）

從實驗結果中，我們觀測到以下現象 ：1）首先，3D-PGT以GPS作為backbone進行預訓練和微調，但相較于GPS的baseline，3D-PGT在8種量子化學性質的預測上平均減少了17.7%的MAE；2）同時，3D-PGT也明顯優于其他2D預訓練方法，以GraphCL為例，其不包含3D結構信息的預訓練所能帶來的性能提升是有限的；3）相較于GraphMVP和3D Infomax兩種前沿的3D預訓練方法，我們的預訓練框架仍然具有性能優勢；4）值得注意的是，當給定準確的3D結構時，3D模型SMP取得了顯著的性能優勢，但當提供給SMP的3D信息來源是粗糙的RDKit計算結果時，3D-PGT在8個預測任務中的6個都擊敗了SMP，體現了3D預訓練范式在效率和性能上優勢。

圖9：不同方法在HOMO預測任務上的推理時間和性能排名

4.2 在其他廣泛下游任務上的泛化能力

除了預測與分子 3D 結構密切相關的量子化學性質外，我們還將下游任務進一步推廣到藥理學、物理學和生物學等領域，且這些任務的輸入只包含 2D 分子圖。我們在 GEOM 數據集中隨機選擇 50 k個具有單個構象的分子進行預訓練，并對 8 個主流下游分子性質分類任務和 6 個回歸任務上進行微調，且這些下游任務包含的數據量都較少，具體結果如下：

圖10：在四個分子性質預測任務（RMSE）和兩個藥物蛋白靶點親和力預測任務（MSE）上的實驗結果

可以看到，3D-PGT 在大多數下游任務上都優于其他 baseline，且與 GPS 相比，3D-PGT 依舊可以預訓練獲得穩定的性能收益。這意味著我們所設計的預訓練框架可以將性能增益推廣到廣泛的下游任務上（除分子性質預測本身外，在兩個藥物和靶蛋白親和力預測任務 Davis 和 KIBA 上也獲得了性能收益），而不局限于和 3D 結構密切相關的量子化學性質預測，且即使在小樣本數據上的微調也能帶來性能增益。

4.3 在大規模數據挑戰賽上的表現

在 OGB-LSC 的 Graph-level 預測任務賽道中，PCQM4Mv2 是一個包含 374 萬分子的大型分子數據集，其中的 337 萬個訓練樣本的 3D 幾何信息通過 DFT 計算得到，且為了接近大規模的虛擬篩選場景，該挑戰沒有提供驗證集和測試集的 3D 構象，且要求使用單個 GPU 在 4 小時內完成 150k 個分子的 HOMO-LUMO gap 的推斷，這意味著在模型的推理階段計算所有測試樣本的幾何結構是不可行的。3D-PGT 在該數據集上的表現與 Leaderboard 排名前列的方法對比結果如下所示：

圖11：在PCQM4Mv@驗證集上的實驗結果

可以看到，與現有的不考慮 3D 結構信息的 GNN 和 Transformer 系列方法相比，3D-PGT 通過引入生成式預訓練任務得到了明顯的性能提升。其中針對GPS 的 baseline，3D-PGT 通過設計的自動預訓練框架降低了 10.6 的 MAE，且在單模型的性能上由于OGB@NeruIPS 2022冠軍方案GPS++。

與此同時，Open Catalyst Challenge 2022是一個致力于利用分子模擬發現高性能催化劑的挑戰賽。該比賽的核心任務是之一利用機器學習模型預測催化劑分子活性。訓練集包含 200 萬個分子，且記錄了每個分子樣本的 3D 結構和總能量。最終，結合 AutoGraph 的算法，3D-PGT 在該挑戰賽上取得了第三的成績，再次印證了該預訓練框架的有效性和競爭力。

圖12：Open Catalyst Challenge 2022 Leaderboard結果

5.  結論

在這項工作中，我們提出了 3D-PGT，一種自動化的 3D 預訓練框架，專注于提升在實 3D 結構不可用的分子性質預測任務上的表現。

3D-PGT 設計了多個生成式預訓練任務，這些任務可以將幾何先驗帶入微調階段。同時，為了更好地自動融合這些預訓練任務并使其效益能夠具有普適性，我們基于分子總能量設計了一個預訓練的 surrogate metric 來搜索每個預任務的自適應權重。本文基于現有主流的量子化學數據集和一系列廣泛的下游任務設計了實驗，且實驗結果表明，3D-PGT 通過預訓練引入潛在的幾何先驗不僅有利于量子化學性質的預測，而且有利于藥理學、物理化學和生物物理學等領域的預測。此外，在 OGB 排行榜上，3D-PGT 在大規模分子預測方面優于頂級解決方案的所有 baseline。

在未來的工作中，考慮到在 NLP 領域 GPT 的預訓練范式的成功，我們會將該框架和 3D 分子預訓練范式推廣到實際的下游應用中，例如開發用于儲存可再生能源的催化劑等。

注：該論文已被KDD 2023接收，如有問題，請聯系wangxu01@4paradigm.com。