DoNews7月12日消息，第四范式基于生成式3D預(yù)訓練大模型在分子性質(zhì)預(yù)測領(lǐng)域的最新研究成果，近日被國際頂會KDD 2023收錄。
論文內(nèi)容已經(jīng)發(fā)布到arxiv，代碼已開源。論文作者是4Paradigm AutoGraph團隊。
分子性質(zhì)預(yù)測是藥物發(fā)現(xiàn)和材料科學等領(lǐng)域中的一個重要問題。考慮到分子的3D結(jié)構(gòu)信息與其性質(zhì)緊密相關(guān)，近年來，一個主流的研究熱點是將分子的3D結(jié)構(gòu)信息與各種以圖學習為代表的機器學習方法結(jié)合，以提高分子性質(zhì)的預(yù)測性能。
然而，由于高昂的計算成本，在大規(guī)模數(shù)據(jù)實時計算分子的3D結(jié)構(gòu)幾乎是不可行的。
圖1：cc(=o)oc1=cc=cc=c1c(=o)o的2D分子圖（左）和3D構(gòu)象圖（右）。
面對這一挑戰(zhàn)，第四范式提出了一個以預(yù)訓練范式為基礎(chǔ)的大模型，專門用于分子性質(zhì)預(yù)測。預(yù)訓練范式隸屬于 AIGC 的一個重要技術(shù)分支[2]，其能處理大規(guī)模數(shù)據(jù)集并提取深層次特征，充分利用了大模型的優(yōu)勢。
本文在分子領(lǐng)域結(jié)合已有的3D分子構(gòu)象進行預(yù)訓練，并在實際下游任務(wù)中只基于分子的2D結(jié)構(gòu)信息進行微調(diào)并進行性質(zhì)預(yù)測，從而在保證高效的前提下進一步提升在下游任務(wù)上的表現(xiàn)。
在這項工作中，第四范式提出了一種自動化3D預(yù)訓練框架：3D-PGT?；诜肿拥幕瘜W鍵長、鍵角和二面角是對應(yīng)于完整分子3D構(gòu)象的三個基本幾何描述符這一事實，第四范式對應(yīng)設(shè)計了三個生成式預(yù)訓練任務(wù)，使得模型通過預(yù)訓練能夠具備編碼3D幾何結(jié)構(gòu)的能力。而為了自動化分配這三個預(yù)訓練任務(wù)的權(quán)重以融合成一個總體的預(yù)訓練目標函數(shù)，第四范式基于分子總能量設(shè)計了一個surrogate metric，從而自動搜索三個預(yù)訓練任務(wù)的權(quán)重分布。
圖2：分子性質(zhì)預(yù)測問題的核心問題示意：如何兼具更快的推理速度和更小的推理誤差。
為了驗證所設(shè)計的預(yù)訓練框架的有效性，第四范式基于已測定 3D 結(jié)構(gòu)的公開分子數(shù)據(jù)集進行了預(yù)訓練，并在 8 個量子化學性質(zhì)預(yù)測任務(wù)和 12 個涉及藥理學、分子化學的下游任務(wù)上進行微調(diào)和實驗驗證。
結(jié)果表明，3D-PGT能通過3D預(yù)訓練帶來明顯的性能增益，且優(yōu)于其他預(yù)訓練 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ì)作為描述分子特征和行為的基本屬性，在藥理學、材料化學等研究領(lǐng)域和具體應(yīng)用都有著重要意義。然而分子性質(zhì)的測定往往需要借助成本高昂的計算方式，例如以密度泛函理論（Density Functional Theory, DFT）為代表的計算化學方法。
近年來，隨著大規(guī)模量子化學計算和高通量實驗的技術(shù)進步，一個新興的工業(yè)界和學術(shù)界熱門的研究方向，是利用具有適當歸納偏置的機器學習方法并結(jié)合已測定性質(zhì)的大量分子數(shù)據(jù)，實現(xiàn)高效的分子性質(zhì)預(yù)測，并應(yīng)用于一系列實際下游應(yīng)用中，例如大規(guī)模藥物分子篩選，合成材料篩選，催化劑設(shè)計等。
圖3：分子性質(zhì)預(yù)測效率對比示意圖。其中基于量子化學計算的DFT計算特定分子性質(zhì)需要數(shù)個小時，而利用機器學習方法只需要遠小于1秒的時間。該圖出自O(shè)GB@NeurIPS 2022 Challenge，預(yù)測分子的HOMO-LUMO能隙。
在過去的幾年中，一種主流的方法是將分子建模為 2D 圖結(jié)構(gòu)，其中以原子作為節(jié)點，而邊作為化學鍵，將整個分子性質(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的量子化學計算方法首先對分子的三維結(jié)構(gòu)進行優(yōu)化以獲得低能構(gòu)象，然后基于構(gòu)象計算特定的分子性質(zhì)。2D模型近似DFT的整個過程，通過輸入的2D分子圖直接預(yù)測分子性質(zhì)；3D模型在2D分子圖的基礎(chǔ)上還需要進一步輸入3D分子構(gòu)象，以獲得更好的預(yù)測性能。
由于分子的 3D 結(jié)構(gòu)反映了原子和官能團之間的相互作用和相對位置，因此 3D 結(jié)構(gòu)對分子性質(zhì)的理解至關(guān)重要?？紤]到這一點，一系列方法針對分子的 3D 結(jié)構(gòu)信息設(shè)計 3D 模型，從而獲得更好的預(yù)測效果。然而，3D 結(jié)構(gòu)的獲取需要依賴 DFT 等量子化學計算手段，這種昂貴的計算成本導致分子的 3D 結(jié)構(gòu)在許多實際的下游任務(wù)中往往是不可獲取的。
結(jié)合上述兩種主流路線各自的優(yōu)缺點，一個最近比較新穎的角度是將包含已測定 3D 結(jié)構(gòu)信息的數(shù)據(jù)集作為預(yù)訓練數(shù)據(jù)集，基于 3D 結(jié)構(gòu)設(shè)計預(yù)訓練任務(wù)，從而讓模型理解 DFT 由 2D 分子圖計算優(yōu)化 3D 結(jié)構(gòu)的過程，并將該先驗信息遷移到下游 2D 分子性質(zhì)預(yù)測任務(wù)上，從而兼具效率和性能。
目前，主流的同期工作例如 GraphMVP[3]?和3D Infomax[4]?都通過對齊分子 2D 視圖和 3D 視圖的圖級表示向量來設(shè)計預(yù)訓練目標函數(shù)，并通過在主流 benchmark 上的實驗證明了 3D 預(yù)訓練方案的有效性。
綜上，本文繼續(xù)沿著分子 3D-Pretraining 的 Pipeline，設(shè)計了三個基礎(chǔ)的生成式預(yù)訓練任務(wù)，并基于分子總能量設(shè)計目標函數(shù)來自動搜索各預(yù)訓練任務(wù)的權(quán)重，從而構(gòu)建多預(yù)訓練任務(wù)的自動融合框架，并在廣泛的下游任務(wù)上獲得了顯著的預(yù)訓練收益。
3.? 本文的方法
3.1??生成式預(yù)訓練任務(wù)的設(shè)計
DFT 優(yōu)化分子 3D 結(jié)構(gòu)的核心是在勢能面上尋找局部最小值[5]，分子能量出于勢能面上局部最小值時的 3D 結(jié)構(gòu)也被稱為分子構(gòu)象。其中，構(gòu)象所包含的 3D 幾何信息可以被以下三個描述符完整地描述：兩個原子之間的化學鍵長、兩個化學鍵之間形成的夾角，以及三個化學鍵形成的二面角。
生成式預(yù)訓練任務(wù)使模型能夠理解分子從 2D 拓撲到 3D 幾何的基于 DFT 的優(yōu)化過程。由于分子幾何是由電子的量子力學行為決定的，因此生成式預(yù)訓練任務(wù)可以通過學習 3D 構(gòu)象的生成來間接學習量子化學性質(zhì)的預(yù)測?；诖?，本文設(shè)計了三個生成式預(yù)訓練任務(wù)來分別生成這鍵長、鍵角、二面角這三個描述符，具體如下所示：
圖5：分子結(jié)構(gòu)和勢能面關(guān)系示意圖。勢能面將分子能量定義為多個坐標軸的函數(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)信息當做解讀信號，從而讓 backbone 具有編碼 3D 結(jié)構(gòu)信息的能力，并將這種能力作為先驗嵌入到模型并遷移到實際的下游任務(wù)中。
同時，考慮到當鄰居數(shù)為|N|時，分子的鍵角和二面角的計算復(fù)雜度呈 Ο(|N|^2 )和Ο(|N|^3 ) 增長，使得大規(guī)模的 3D 預(yù)訓練實現(xiàn)成本過高。本文基于 RGC（Runtime Geometry Calculation）重新設(shè)計了鍵角和二面角的目標函數(shù)，用每個原子的所涉及的鍵角和與每個化學鍵設(shè)計的二面角的和來代替所有鍵角和二面角的預(yù)測，從而將計算復(fù)雜度降低到線性級別。
3.2??自動化多預(yù)訓練任務(wù)融合框架
由于第四范式同時定義了多個預(yù)訓練任務(wù)，且各預(yù)訓練任務(wù)的損失函數(shù)對于模型參數(shù)的梯度優(yōu)化方向并不一致，因此各自預(yù)訓練任務(wù)的權(quán)重分配是一個需要考慮的問題。如何自動化且最更好地分配多預(yù)訓練任務(wù)的權(quán)重，以預(yù)期在下游任務(wù)上獲得更好的性能增益，是設(shè)計多預(yù)訓練任務(wù)融合框架的核心。
由于鍵長、鍵角和二面角都是局部描述符，因此這里需要一個明確的 surrogate metric 來評估預(yù)訓練后的 backbone 對于分子 3D 結(jié)構(gòu)整體的編碼能力而不是局部編碼能力。
對此，本文考慮得到分子總能量和分子 3D 結(jié)構(gòu)之間的對應(yīng)關(guān)系，基于分子總能量設(shè)計了一個目標函數(shù)，并基于此設(shè)計了 bi-level 的油畫框架來搜索三個生成式預(yù)訓練任務(wù)各自的權(quán)重。
最終，第四范式整體預(yù)訓練的 pipeline 如下圖所示。在預(yù)訓練階段，第四范式首先基于預(yù)訓練分子的低能構(gòu)象設(shè)計了三個生成式預(yù)訓練任務(wù)，并基于分子低能構(gòu)象對應(yīng)的分子總能量設(shè)計了一個 surrogate metric 來搜索三個預(yù)訓練任務(wù)的各自權(quán)重；而在下游的微調(diào)階段，由于我們已經(jīng)在預(yù)訓練階段引入了分子幾何先驗，因此通過微調(diào)即可在實際下游任務(wù)中獲得性能增益。
圖6：3D-PGT的預(yù)訓練框架示意圖
4.? 實驗結(jié)果
在本文中，第四范式針對性的設(shè)計了一系列實驗來證明 3D 預(yù)訓練這一技術(shù)路線和本文設(shè)計的預(yù)訓練任務(wù)及自動融合框架的有效性，從而面臨當只有 2D 分子圖可以用于預(yù)測時，可以避免為每個分子生成 3D 構(gòu)象的巨大計算成本，并能通過微調(diào)獲得明顯的性能收益。
本文主要在 3 個包含 3D 結(jié)構(gòu)信息的數(shù)據(jù)集上進行預(yù)訓練，并主要在 12 個下游任務(wù)上進行微調(diào)和效果驗證，數(shù)據(jù)集的詳細統(tǒng)計數(shù)據(jù)如下所示：
圖7：數(shù)據(jù)集的統(tǒng)計細節(jié)
4.1 量子化學性質(zhì)預(yù)測的表現(xiàn)
本文首先在流行的基準數(shù)據(jù)集QM9上評估了3D-PGT的量子化學領(lǐng)域的性質(zhì)預(yù)測能力。第四范式首先從QM9包含的134k單一分子構(gòu)象樣本中隨機選取50k個攜帶構(gòu)象信息的分子樣本用于預(yù)訓練，并從剩余的樣本中選取50k個分子樣本并屏蔽其3D結(jié)構(gòu)信息用于微調(diào)和評測。結(jié)果如下：
圖8：QM9數(shù)據(jù)集實驗對比結(jié)果，評價指標為MAE（平均絕對誤差）
從實驗結(jié)果中，第四范式觀測到以下現(xiàn)象 ：1）首先，3D-PGT以GPS作為backbone進行預(yù)訓練和微調(diào)，但相較于GPS的baseline，3D-PGT在8種量子化學性質(zhì)的預(yù)測上平均減少了17.7%的MAE；2）同時，3D-PGT也明顯優(yōu)于其他2D預(yù)訓練方法，以GraphCL為例，其不包含3D結(jié)構(gòu)信息的預(yù)訓練所能帶來的性能提升是有限的；3）相較于GraphMVP和3D Infomax兩種前沿的3D預(yù)訓練方法，我們的預(yù)訓練框架仍然具有性能優(yōu)勢；4）值得注意的是，當給定準確的3D結(jié)構(gòu)時，3D模型SMP取得了顯著的性能優(yōu)勢，但當提供給SMP的3D信息來源是粗糙的RDKit計算結(jié)果時，3D-PGT在8個預(yù)測任務(wù)中的6個都擊敗了SMP，體現(xiàn)了3D預(yù)訓練范式在效率和性能上優(yōu)勢。
圖9：不同方法在HOMO預(yù)測任務(wù)上的推理時間和性能排名
4.2 在其他廣泛下游任務(wù)上的泛化能力
除了預(yù)測與分子 3D 結(jié)構(gòu)密切相關(guān)的量子化學性質(zhì)外，第四范式還將下游任務(wù)進一步推廣到藥理學、物理學和生物學等領(lǐng)域，且這些任務(wù)的輸入只包含 2D 分子圖。
第四范式在 GEOM 數(shù)據(jù)集中隨機選擇 50 k個具有單個構(gòu)象的分子進行預(yù)訓練，并對 8 個主流下游分子性質(zhì)分類任務(wù)和 6 個回歸任務(wù)上進行微調(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ù)訓練獲得穩(wěn)定的性能收益。這意味著第四范式所設(shè)計的預(yù)訓練框架可以將性能增益推廣到廣泛的下游任務(wù)上（除分子性質(zhì)預(yù)測本身外，在兩個藥物和靶蛋白親和力預(yù)測任務(wù) Davis 和 KIBA 上也獲得了性能收益），而不局限于和 3D 結(jié)構(gòu)密切相關(guān)的量子化學性質(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 萬個訓練樣本的 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ù)訓練任務(wù)得到了明顯的性能提升。其中針對GPS 的 baseline，3D-PGT 通過設(shè)計的自動預(yù)訓練框架降低了 10.6 的 MAE，且在單模型的性能上由于OGB@NeruIPS 2022冠軍方案GPS++。
與此同時，Open Catalyst Challenge 2022是一個致力于利用分子模擬發(fā)現(xiàn)高性能催化劑的挑戰(zhàn)賽。該比賽的核心任務(wù)是之一利用機器學習模型預(yù)測催化劑分子活性。
訓練集包含 200 萬個分子，且記錄了每個分子樣本的 3D 結(jié)構(gòu)和總能量。最終，結(jié)合 AutoGraph 的算法，3D-PGT 在該挑戰(zhàn)賽上取得了第三的成績，再次印證了該預(yù)訓練框架的有效性和競爭力。
圖12：Open Catalyst Challenge 2022 Leaderboard結(jié)果
5.? 結(jié)論
在這項工作中，第四范式提出了 3D-PGT，一種自動化的 3D 預(yù)訓練框架，專注于提升在實 3D 結(jié)構(gòu)不可用的分子性質(zhì)預(yù)測任務(wù)上的表現(xiàn)。
3D-PGT 設(shè)計了多個生成式預(yù)訓練任務(wù)，這些任務(wù)可以將幾何先驗帶入微調(diào)階段。同時，為了更好地自動融合這些預(yù)訓練任務(wù)并使其效益能夠具有普適性，第四范式基于分子總能量設(shè)計了一個預(yù)訓練的 surrogate metric 來搜索每個預(yù)任務(wù)的自適應(yīng)權(quán)重。
本文基于現(xiàn)有主流的量子化學數(shù)據(jù)集和一系列廣泛的下游任務(wù)設(shè)計了實驗，且實驗結(jié)果表明，3D-PGT 通過預(yù)訓練引入潛在的幾何先驗不僅有利于量子化學性質(zhì)的預(yù)測，而且有利于藥理學、物理化學和生物物理學等領(lǐng)域的預(yù)測。此外，在 OGB 排行榜上，3D-PGT 在大規(guī)模分子預(yù)測方面優(yōu)于頂級解決方案的所有 baseline。
在未來的工作中，考慮到在 NLP 領(lǐng)域 GPT 的預(yù)訓練范式的成功，我們會將該框架和 3D 分子預(yù)訓練范式推廣到實際的下游應(yīng)用中，例如開發(fā)用于儲存可再生能源的催化劑等。
