在刚刚落幕的由 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
原创文章,作者:机器之心,如若转载,请注明出处:https://www.iaiol.com/news/28725