編輯 | KX
Uni-Mol 是深勢科技於 2022 年釋出的一款基於分子三維結構的通用大模型，其效能優越、模型泛化能力強，在小分子性質預測、蛋白靶點預測、量子化學性質預測、MOF 材料吸附效能預測等任務上都超越了現有的解決方案。
今年 3 月，深勢科技與清華大學等合作，提出基於 Uni-Mol 的領域專屬模型 Uni-MOF，用於預測各類工況下奈米多孔材料對各類氣體的吸附效能，預測精度高達 0.98。
近日，深勢科技與北京大學合作，推出新一代模型 Uni-Mol+，迭代升級後的模型擁有更大的引數量，更多的預訓練資料量，並展現出更強大的通用性。
Uni-Mol+ 是一種利用 3D 構象進行精確量子化學屬性預測的深度學習方法。基準測試結果表明，Uni-Mol+ 顯著提高了各種資料集中 QC 屬性預測的準確性。
相關研究以「Data-driven quantum chemical property prediction leveraging 3D conformations with Uni-Mol+」為題，於 8 月 19 日釋出在《Nature Communications》上。
量子化學 (QC) 性質預測對於計算材料和藥物設計至關重要，但依賴於昂貴的電子結構計算，如密度泛函理論 (DFT)。
深度學習方法使用 1D SMILES 或 2D 圖作為輸入來加速這一過程，但難以實現高精度，因為大多數 QC 性質都依賴於精細的 3D 分子平衡構象。
為了應對這一挑戰，深勢科技提出了一種名為 Uni-Mol+ 的方法。
Uni-Mol+ 憑藉精心設計的模型主幹和訓練策略，在各項基準測試中表現出優異的效能。
研究的主要貢獻可以概括如下：
研究人員利用從 RDKit 生成的構象到 DFT 平衡構象的構象最佳化，開發了一種用於 QC 屬性預測的新範例。 透過生成偽軌跡（pseudo trajectory）並從中取樣策略，基於伯努利分佈和均勻分佈的混合，建立了一種用於 3D 構象最佳化的新訓練策略。 Uni-Mol+ 的整個框架具有重要的經驗價值，因為它在兩個廣受認可的基準 PCQM4MV2 和 Open Catalyst 2020 (OC20) 上的效能明顯優於之前的研究。
Uni-Mol+ 概述
對於任何分子，Uni-Mol+ 首先透過廉價方法（例如來自 RDKit 和 OpenBabel 的基於模板的方法）獲得原始 3D 構象。然後，它透過原始構象的迭代更新過程學習目標構象，即由 DFT 最佳化的平衡構象。在最後一步中，根據學習到的構象預測 QC 屬性。
為了有效地學習這個構象更新過程，研究人員提出了一個雙軌 Transformer 模型主幹和一種新穎的訓練方法。
Uni-Mol+ 的模型主幹是一個雙軌 Transformer，由一個原子表示軌道和一個對錶示軌道組成。
與之前 Uni-Mol 中使用的 Transformer 主幹相比，進行了兩項重大更新：
（1）透過原子表示的外積（稱為 OuterProduct）增強對錶示，進行原子到對的通訊，並使用三角運算元（稱為 TriangularUpdate）來增強 3D 幾何資訊。這兩個運算元在 AlphaFold2 中被證明是有效的。
（2）採用迭代過程不斷更新 3D 座標以達到平衡構象。使用 R 表示構象最佳化的 rounds 數。
為了學習構象更新過程，研究人員提出了一種新穎的訓練策略。從 RDKit 生成的原始構象和 DFT 平衡構象之間的軌跡中取樣構象，並使用取樣的構象作為輸入來預測平衡構象。必須注意的是，在許多資料集中，實際軌跡通常是未知的；因此，研究人員使用一種假定兩個構象之間存線上性過程的偽軌跡。
此外，還設計了一種取樣策略，用於從偽軌跡中獲取構象，作為模型在訓練期間的輸入。該策略混合使用伯努利分佈和均勻分佈。伯努利分佈解決了 (1) 訓練和推理之間的分佈轉變，以及 (2) 增強了從平衡構象到 QC 屬性的精確對映的學習。同時，均勻分佈生成額外的中間狀態作為模型輸入，有效地增強了輸入構象。
基準測試
研究人員在兩個大規模資料集基準 PCQM4MV2 和 Open Catalyst 2020 (OC20) 上評估了 Uni-Mol+ 的效能。
首先，將之前提交給 PCQM4MV2 排行榜的模型作為基準。除了預設的 12 層模型外，研究人員還評估了 Uni-Mol+ 的效能，其兩個變體分別由 6 層和 18 層組成。這旨在探索當模型引數大小改變時模型效能如何變化。
結果如下：
（1）Uni-Mol+ 在單模型效能驗證資料上比之前的 SOTA 高出 0.0079，相對提高了 11.4%。
（2）Uni-Mol+ 的所有三種變體都比之前的基線表現出顯著的效能提升。（3）儘管 6 層的 Uni-Mol+ 的模型引數少得多，但它的表現優於所有之前的基線。
（4）將層數從 6 層增加到 12 層可顯著提高準確度，並以相當大的優勢超越所有基線。
（5）18 層的 Uni-Mol+ 表現出最高的效能，以顯著的優勢超越所有基線。這些發現強調了 Uni-Mol+ 的有效性。
（6）單個 18 層 Uni-Mol+ 模型在排行榜（測試開發集）上的表現值得關注，特別是因為它超越了之前最先進的方法，而無需使用整合或其他技術。相比之下，之前最先進的 GPS++ 依賴於 112 個模型整合，幷包括驗證集進行訓練。
Open Catalyst 2020 (OC20) 資料集專門用於促進催化劑發現和最佳化的機器學習模型的開發。在該研究中，重點關注始結構到鬆弛能量（IS2RE）任務。
研究人員對 OC20 IS2RE 驗證和測試集上的各種模型進行了效能比較，如表 2 所示。從表中可以看出，Uni-Mol+ 在平均絕對誤差 (MAE) 和閾值內能量 (EwT) 方面都明顯優於所有之前的基線。這證明了 Uni-Mol+ 的卓越效能。研究結果強調了 Uni-Mol+ 在捕捉材料系統中複雜相互作用方面的有效性，以及它在各種計算材料科學任務中廣泛應用的潛力。
消融研究
研究人員將對 Uni-Mol+ 進行了全面的消融研究。對 PCQM4Mv2 資料集進行了消融研究，採用預設的 12 層 Uni-Mol+ 配置。研究結果總結在表 3 中，其中 No.1 是預設設定，No.2–7 重點檢查模型主幹，No. 8–No. 17 重點檢查訓練策略。
研究結果如下：
（1）比較 No. 8、No. 9 和 No. 10，發現僅從一種構象中取樣效果不佳。
（2）透過比較 No. 8、No. 9 和 No. 11，可以推斷出從 RDKit 和目標構象的混合中取樣會產生令人滿意的結果（有效 MAE 為 0.0697）。但是，如果僅從目標和中間構象（No. 12）取樣，結果並不令人滿意（有效 MAE 為 0.0753）。這一結果表明，從 w1.0 中取樣是必要的，因為它減少了訓練和推理之間的分佈偏移。
（3）從三種構象型別中取樣的預設策略（No. 1）表現出最佳效能。
（4）改變混合分佈的權重（No. 13–17）不會導致比預設策略更好的效能。此外，隨著 w0.0 的減少，效能會變差。這表明預設加權方案適合這項任務。
（5）比較 No.18 和 No.1 的結果後，很明顯，Noisy Nodes（No.18，有效 MAE 為 0.0760）的效能明顯低於 Uni-Mol+（No.1，有效 MAE 為 0.0696）。這種巨大的效能差距（0.0760 vs. 0.0696），凸顯了所提出的訓練策略比以前採用的策略更高效。
（6）對比 No.19 和 No.18，發現在使用噪聲節點策略時，之前研究中採用的模型結構比使用 Uni-Mol+ 的主幹結構產生的結果更差。這一發現進一步證明了 Uni-Mol+ 的主幹結構優於之前提出的模型架構。
總之，消融研究證明了 Uni-Mol+ 中採用的預設取樣策略的有效性，強調了利用不同構象混合物來實現卓越效能的重要性。
構象學習的視覺化分析
除了 QC 性質預測外，Uni-Mol+ 還可以預測平衡構象。雖然該研究主要集中在 QC 屬性預測上，並且證明了 Uni-Mol+ 的有效性，但視覺化的結果可以幫助更好地理解 Uni-Mol+ 的工作原理。因此，研究人員還為 PCQM4MV2 資料集中 Uni-Mol+ 的構象學習提供了兩個額外的分析。
第一個分析評估預測的構象。如圖 2 所示，Uni-Mol+ 可以有效地預測平衡構象。此外，隨著更新迭代次數的增加，RMSD 變小，進一步證明了所提出的迭代座標更新的有效性。
第二個分析旨在證明 Uni-Mol+ 可以預測較低能量的構象，接近平衡構象。如圖 3 所示，Uni-Mol+ 可以預測能量較低的構象。此外，初始構象和預測構象之間的能量差分佈，與初始構象和平衡構象之間的能量差分佈密切一致。這種相似性證明了 Uni-Mol+ 在準確預測平衡構象方面的有效性。
上述結果為所提出的 Uni-Mol+ 的有效性提供了額外的證據，因為它確實可以預測較低能量的構象，並迭代接近目標 DFT 構象。
總之，該研究提出了一種新穎的方法，能夠透過輔助任務——構象最佳化，準確預測量子化學性質。這種方法有望提高高通量篩選的效率，並促進創新材料和分子設計。