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

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

一、背景

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

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

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

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

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

二、技術(shù)

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

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

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

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

數(shù)據(jù)來源: https://eval.ai/web/challenges/challenge-page/712/leaderboard/1950

在 GeoEnsemble 框架中，為了對原子之間復雜的動態(tài)交互關(guān)系進行建模，我們在 ICLR 2022 發(fā)表的圖動力學神經(jīng)網(wǎng)絡（GMN）[1]的基礎(chǔ)上進行了改進，提出了 GMN-OC 模型。

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

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

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

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

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

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

全局共享表示模塊示意圖

整個 GMN-OC 模型結(jié)構(gòu)如下圖。同時，我們在計算時，使用了 Multi-head 的顯存優(yōu)化方法，使得模型可以更好的應對大數(shù)據(jù)的處理。

GMN-OC 模型整體架構(gòu)

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

GeoEnsemble 在訓練上的改進

三、展望

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

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

我們將在近期開源這次比賽使用的模型和訓練推理源碼，以助力人工智能在量子化學模擬和電子結(jié)構(gòu)計算等基礎(chǔ)研究領(lǐng)域的應用。在未來，騰訊 AI Lab 將持續(xù)研發(fā)和落地 AI 新技術(shù)，推動 AI 在交叉學科中的新應用，探索 AI 賦能科學發(fā)現(xiàn)的新范式。

比賽鏈接：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