當AI遇上量子化學，這是NeurIPS 2022挑戰賽的冠軍解決方案

在剛剛落幕的由 Meta AI 研究院及卡耐基梅隆大學（CMU）聯合機器學習頂級會議 NeurIPS 共同舉辦的第二屆 Open Catalyst Challenge (OCP)競賽中，由騰訊 AI Lab 領頭，中國人民大學，清華大學以及香港中文大學組成的聯合團隊 TTRC 以 0.396eV 絕對誤差的成績獲得第一，相對于去年的最好成績，提升達到 27.6%。

一、背景

Open Catalyst Project (OCP)是由 Meta AI 和卡耐基梅隆大學聯合發起的一個科研項目，旨在使用人工智能算法加速可用于再生能源存儲的催化劑的發現。其核心目的是通過尋找高效且經濟的催化劑來解決再生能源存儲的問題。

在這一領域，傳統方法是基于量子力學（密度泛函理論 DFT）的模擬計算催化劑表面和目標吸附物的結合能來測試和評估新的催化劑結構。然而，這類模擬的一大缺點是其巨大的計算成本，通常單個組合模擬就需要耗費 24 小時以上。這使得我們無法高效且大量地篩選潛在的催化劑結構。因此，如何利用機器學習算法去建模量子化學模型進而實現對于催化劑結構的有效篩選還是一個非常具有前景的方向。

傳統基于量子力學的計算流程（引自官方資料）

給定一個催化劑 - 吸附物組成的反應系統，經典的量子力學模擬計算分為兩個步驟，第一步是基于密度泛函理論計算系統中微觀粒子的等效受力。第二步則是基于這個受力迭代更新系統中原子核的位置。直到受力收斂。這樣就可以得到這個系統低能態，即松弛狀態，進而計算催化劑和吸附物的結合能。

本次比賽參賽主題則是構造機器學習模型預測一個由催化劑 - 吸附物組成的反應系統的松弛狀態能量，即 IS2RE(Initial state to relaxed energy)。同上一屆不同的是，除了提供催化劑 - 吸附物以及對應能量標簽（46 萬）外，主辦方還額外提供了 200 萬的由密度泛函計算得到的靜態反應系統的坐標及其對應的力和能量。基于這些數據，模型可以去預測靜態結構下的等效受力，即 S2EF(Structure to energy and force)，使得模型可以更好探索量子力學計算的部分建模。

二、技術

自從 OCP 項目發起以來，由于其重要的科學意義和龐大的數據規模以及具有挑戰的任務，吸引了 DeepMind，MSRA，達摩院，MILA，MIT，上海交大等機構的關注和參與。在今年的第二屆比賽中，我們提出的 GeoEnsemble 框架以 0.396eV 絕對誤差的成績獲得第一，相對于去年比賽最好成績 Graphormer，提升達到 27.6%。

非公開測試集排行榜，帶下劃線的為官方 Baseline。（引自官方資料）

本屆比賽的方案相對于上一屆方案的提升。（引自官方資料）

限定比賽用訓練集的方法，在 Public Leaderboard 性能表現。

數據來源: https://eval.ai/web/challenges/challenge-page/712/leaderboard/1950

在 GeoEnsemble 框架中，為了對原子之間復雜的動態交互關系進行建模，我們在 ICLR 2022 發表的圖動力學神經網絡（GMN）[1]的基礎上進行了改進，提出了 GMN-OC 模型。

GMN-OC 模型的輸入是一個由原子構成的幾何圖，幾何圖中包含了幾何特征（原子的三維坐標）和非幾何特征（原子的類型），模型可以預測輸出幾何向量（原子的受力）和非幾何標量（系統能量）。

GMN-OC 整體輸入和輸出流程

在 GMN-OC 模型中，我們構造了一個基于多通道的幾何特征 O(3)等變函數 與不變函數  用來處理幾何特征和非幾何特征的交互。

基于這兩個函數，我們構建了一個基于消息傳遞的圖神經網絡，在 GMN-OC 的每一個網絡層中，會基于 O(3)等變與不變函數進行消息傳遞與聚合，從而實現對每個原子的幾何特征和非幾何特征進行更新。

基于消息傳播的幾何特征更新過程

在這個基礎模型上，我們進一步引入了一個全局共享表示模塊（Global Representation Module）建模在 DFT 計算中可能涉及到的粒子之間的全局交互信息。同時也保持了模型的等變性。

全局共享表示模塊示意圖

整個 GMN-OC 模型結構如下圖。同時，我們在計算時，使用了 Multi-head 的顯存優化方法，使得模型可以更好的應對大數據的處理。

GMN-OC 模型整體架構

此外在訓練技術上，為了充分利用賽事提供的兩組訓練數據，我們使用了 Pretrain-Finetune、Multi-task Learning 等優化技術，進一步提高模型預測精度。此外，我們結合 GMN-OC/SCN/GemNet 等多個模型，構建了多模型融合方案 GeoEnsemble。為了提高模型訓練速度，我們還使用了混合精度訓練等技術，并在大規模集群上完成了分布式訓練。

GeoEnsemble 在訓練上的改進

三、展望

利用人工智能技術助力自然科學領域中的探索和發現，已經成為近年來人工智能領域備受矚目的應用方向之一。得益于高性能計算能力和前沿人工智能技術的不斷發展，基礎科學領域的研究者得以利用人工智能算法去加速相關領域中復雜、大規模的計算和模擬任務，如蛋白質折疊、小分子結合能計算、催化劑發現等。

騰訊 AI Lab，基于在人工智能算法研究中的積累，在生命科學，物理建模等多個應用方向上取得了重要的突破。例如首個大規模小分子預訓練模型 GROVER [2]及骨架躍遷生成模型 [3]，自研蛋白質折疊模型 tFold [3] 和基于序列預訓練的抗體結構預測模型 tFold-Ab [4]，基于等變性的蛋白質對接模型 EquiDock [5]，基于圖動力學網絡的蛋白質動態結構預測模型 EGHN [6]等。同時，騰訊云深 AI ****物發現平臺研發的基于等變圖神經網絡的分子能量框架 DeepQC，可以實現對類****分子的高精度的量化計算。

我們將在近期開源這次比賽使用的模型和訓練推理源碼，以助力人工智能在量子化學模擬和電子結構計算等基礎研究領域的應用。在未來，騰訊 AI Lab 將持續研發和落地 AI 新技術，推動 AI 在交叉學科中的新應用，探索 AI 賦能科學發現的新范式。

比賽鏈接：https://opencatalystproject.org/challenge.html

項目主頁：https://ai.tencent.com/ailab/ml/ocp/index.html

引用：

1. Equivariant Graph Mechanics Networks with Constraints

2. Self-Supervised Graph Transformer on Large-Scale Molecular Data

3. A novel scalarized scaffold hopping algorithm with graph-based variational autoencoder for discovery of JAK1 inhibitors

4. When homologous sequences meet structural decoys: Accurate contact prediction by tFold in CASP14—(tFold for CASP14 contact prediction)

5. tFold-Ab: Fast and Accurate Antibody Structure Prediction without Sequence Homologs

6. Independent SE(3)-Equivariant Models for End-to-End Rigid Protein Docking

7. Equivariant Graph Hierarchy-Based Neural Networks