近日，第四范式基于生成式3D預(yù)訓(xùn)練大模型在分子性質(zhì)預(yù)測領(lǐng)域的最新研究成果，被國際頂會KDD 2023收錄。

下文為論文內(nèi)容，作者是4Paradigm AutoGraph團(tuán)隊。

論文標(biāo)題：Automated 3D Pre-Training for Molecular Property Prediction

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

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

1.  簡介

分子性質(zhì)預(yù)測是藥物發(fā)現(xiàn)和材料科學(xué)等領(lǐng)域中的一個重要問題。考慮到分子的3D結(jié)構(gòu)信息與其性質(zhì)緊密相關(guān)，近年來，一個主流的研究熱點是將分子的3D結(jié)構(gòu)信息與各種以圖學(xué)習(xí)為代表的機(jī)器學(xué)習(xí)方法結(jié)合，以提高分子性質(zhì)的預(yù)測性能。然而，由于高昂的計算成本，在大規(guī)模數(shù)據(jù)實時計算分子的3D結(jié)構(gòu)幾乎是不可行的[1]。

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

面對這一挑戰(zhàn)，我們提出了一個以預(yù)訓(xùn)練范式為基礎(chǔ)的大模型，專門用于分子性質(zhì)預(yù)測。預(yù)訓(xùn)練范式隸屬于 AIGC 的一個重要技術(shù)分支[2]，其能處理大規(guī)模數(shù)據(jù)集并提取深層次特征，充分利用了大模型的優(yōu)勢。本文在分子領(lǐng)域結(jié)合已有的3D分子構(gòu)象進(jìn)行預(yù)訓(xùn)練，并在實際下游任務(wù)中只基于分子的2D結(jié)構(gòu)信息進(jìn)行微調(diào)并進(jìn)行性質(zhì)預(yù)測，從而在保證高效的前提下進(jìn)一步提升在下游任務(wù)上的表現(xiàn)。

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

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

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

2.  背景介紹

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

圖3：分子性質(zhì)預(yù)測效率對比示意圖。其中基于量子化學(xué)計算的DFT計算特定分子性質(zhì)需要數(shù)個小時，而利用機(jī)器學(xué)習(xí)方法只需要遠(yuǎn)小于1秒的時間。該圖出自O(shè)GB@NeurIPS 2022 Challenge，預(yù)測分子的HOMO-LUMO能隙。

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

圖4：現(xiàn)有分子性質(zhì)預(yù)測方法的一般流程。基于DFT的量子化學(xué)計算方法首先對分子的三維結(jié)構(gòu)進(jìn)行優(yōu)化以獲得低能構(gòu)象，然后基于構(gòu)象計算特定的分子性質(zhì)。2D模型近似DFT的整個過程，通過輸入的2D分子圖直接預(yù)測分子性質(zhì)；3D模型在2D分子圖的基礎(chǔ)上還需要進(jìn)一步輸入3D分子構(gòu)象，以獲得更好的預(yù)測性能。

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

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

綜上，本文繼續(xù)沿著分子 3D-Pretraining 的 Pipeline，設(shè)計了三個基礎(chǔ)的生成式預(yù)訓(xùn)練任務(wù)，并基于分子總能量設(shè)計目標(biāo)函數(shù)來自動搜索各預(yù)訓(xùn)練任務(wù)的權(quán)重，從而構(gòu)建多預(yù)訓(xùn)練任務(wù)的自動融合框架，并在廣泛的下游任務(wù)上獲得了顯著的預(yù)訓(xùn)練收益。

3.  本文的方法

3.1  生成式預(yù)訓(xùn)練任務(wù)的設(shè)計

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

生成式預(yù)訓(xùn)練任務(wù)使模型能夠理解分子從 2D 拓?fù)涞?3D 幾何的基于 DFT 的優(yōu)化過程。由于分子幾何是由電子的量子力學(xué)行為決定的，因此生成式預(yù)訓(xùn)練任務(wù)可以通過學(xué)習(xí) 3D 構(gòu)象的生成來間接學(xué)習(xí)量子化學(xué)性質(zhì)的預(yù)測。基于此，本文設(shè)計了三個生成式預(yù)訓(xùn)練任務(wù)來分別生成這鍵長、鍵角、二面角這三個描述符，具體如下所示：

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

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

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

3.2  自動化多預(yù)訓(xùn)練任務(wù)融合框架

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

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

圖6：3D-PGT的預(yù)訓(xùn)練框架示意圖

4.  實驗結(jié)果

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

本文主要在 3 個包含 3D 結(jié)構(gòu)信息的數(shù)據(jù)集上進(jìn)行預(yù)訓(xùn)練，并主要在 12 個下游任務(wù)上進(jìn)行微調(diào)和效果驗證，數(shù)據(jù)集的詳細(xì)統(tǒng)計數(shù)據(jù)如下所示：

圖7：數(shù)據(jù)集的統(tǒng)計細(xì)節(jié)

4.1 量子化學(xué)性質(zhì)預(yù)測的表現(xiàn)

本文首先在流行的基準(zhǔn)數(shù)據(jù)集QM9上評估了3D-PGT的量子化學(xué)領(lǐng)域的性質(zhì)預(yù)測能力。我們首先從QM9包含的134k單一分子構(gòu)象樣本中隨機(jī)選取50k個攜帶構(gòu)象信息的分子樣本用于預(yù)訓(xùn)練，并從剩余的樣本中選取50k個分子樣本并屏蔽其3D結(jié)構(gòu)信息用于微調(diào)和評測。結(jié)果如下：

圖8：QM9數(shù)據(jù)集實驗對比結(jié)果，評價指標(biāo)為MAE（平均絕對誤差）

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

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

4.2 在其他廣泛下游任務(wù)上的泛化能力

除了預(yù)測與分子 3D 結(jié)構(gòu)密切相關(guān)的量子化學(xué)性質(zhì)外，我們還將下游任務(wù)進(jìn)一步推廣到藥理學(xué)、物理學(xué)和生物學(xué)等領(lǐng)域，且這些任務(wù)的輸入只包含 2D 分子圖。我們在 GEOM 數(shù)據(jù)集中隨機(jī)選擇 50 k個具有單個構(gòu)象的分子進(jìn)行預(yù)訓(xùn)練，并對 8 個主流下游分子性質(zhì)分類任務(wù)和 6 個回歸任務(wù)上進(jìn)行微調(diào)，且這些下游任務(wù)包含的數(shù)據(jù)量都較少，具體結(jié)果如下：

圖10：在四個分子性質(zhì)預(yù)測任務(wù)（RMSE）和兩個藥物蛋白靶點親和力預(yù)測任務(wù)（MSE）上的實驗結(jié)果

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

4.3 在大規(guī)模數(shù)據(jù)挑戰(zhàn)賽上的表現(xiàn)

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

圖11：在PCQM4Mv@驗證集上的實驗結(jié)果

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

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

圖12：Open Catalyst Challenge 2022 Leaderboard結(jié)果

5.  結(jié)論

在這項工作中，我們提出了 3D-PGT，一種自動化的 3D 預(yù)訓(xùn)練框架，專注于提升在實 3D 結(jié)構(gòu)不可用的分子性質(zhì)預(yù)測任務(wù)上的表現(xiàn)。

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

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

參考文獻(xiàn)

[1] Learning neural generative dynamics for molecular conformation generation. Arxiv 2021

[2]  Language models are few-shot learners. NeurIPS 2020

[3] Pre-training Molecular Graph Representation with 3D Geometry. ICLR 2022

[4] 3D Infomax improves GNNs for Molecular Property Prediction. ICML 2022

[5] GEOM, energy-annotated molecular conformations for property prediction and molecular generation. Scientific Data 2022