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

一、背景

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

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

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

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

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

二、技術(shù)

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

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

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

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

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

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

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

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

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

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

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

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

全局共享表示模塊示意圖

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

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

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

GeoEnsemble 在訓(xùn)練上的改進(jìn)

三、展望

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

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

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

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

項(xiàng)目主頁：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