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