zl程序教程

您现在的位置是:首页 >  其它

当前栏目

Uni-Mol+:Highly Accurate Quantum Chemical PropertyPrediction with Uni-Mol+

with uni
2023-09-14 09:14:41 时间

Uni-Mol+:高精度的量子化学性质预测

Uni-Mol+ 预印论文:https://github.com/dptech-corp/Uni-Mol/blob/main/unimol_plus/paper/unimol_plus_preprint.pdf

 Uni-Mol+ 代码:Uni-Mol/README.md at main · dptech-corp/Uni-Mol · GitHub

最近,我们对Uni-Mol的模型框架进行了大幅升级,推出了Uni-Mol+,并在国际权威学术竞赛OGB-LSC的量子化学性质预测上获得了榜首。 


基础概念: 

1、DFT-generated 3D conformation是什么

DFT-generated 3D conformation指的是通过密度泛函理论(DFT)计算得到的分子的三维构象。

DFT计算得到的分子构象可以用来理解分子的化学性质、反应机理、分子间相互作用等。在DFT计算中,需要先确定分子的初始构象,然后通过计算得到优化后的最稳定构象。DFT可以考虑分子中电子的相互作用,因此DFT-generated 3D conformation相对于其他构象预测方法来说,具有更高的精度和可靠性。DFT-generated 3D conformation也可以用于分子模拟、分子对接、药物设计等领域。

2、HOMO-LUMO gap label是什么?

HOMO-LUMO gap(Highest Occupied Molecular Orbital-Lowest Unoccupied Molecular Orbital gap)指的是一个分子中最高占据分子轨道(HOMO)和最低未占据分子轨道(LUMO)之间的能量差。

HOMO-LUMO gap通常被用来描述分子的化学性质和光学性质,例如电子亲和性、电离势、吸收光谱、发射光谱等。HOMO-LUMO gap越小,表示分子越容易接受或者捐赠电子,也越容易被激发并吸收或发射光子。因此,HOMO-LUMO gap对于理解分子的化学性质和光学性质非常重要。

3、LTI 是什么?

LTI是指线性和时不变(Linear and time invariance)系统,即具有齐次性(Homogeneity)、可叠加性(Superposition)和时不变性(Time invariance)三个特征的系统。

而在论文中提到的LTI策略是指线性轨迹注入(linear trajectory injection)策略,这是一种用于优化构象的方法,LTI方法提供了多个分子的输入构象,可以作为有效的数据增强策略


一、数据集Setting

OGB-LSC (Open Graph Benchmark, Large Scale Challenge) 是一项由斯坦福大学发起的学术竞赛,旨在评估机器学习在大规模图数据上的表现。该竞赛首次在KDD CUP 2021上举办,吸引了来自DeepMind、微软、NVIDIA、UCLA等顶尖企业和高校的500多个参赛队伍,备受业界关注。近年来,越来越多的新型图机器学习模型也加入到这个比赛中,以证明自己的模型性能。可以说,OGB-LSC已成为公认的检验图机器学习模型性能的最佳试金石,类似于ImageNet在图像领域的地位。 

OGB-LSC代表Open Graph Benchmark中的Large-Scale Challenge(大规模挑战),是一个基准测试任务,旨在对图神经网络的性能进行评估和比较。OGB-LSC基于来自真实世界的大规模图数据集,包括了节点分类、边预测和图分类三个任务。该任务旨在推动图神经网络的发展,提高图数据处理任务的准确性和效率。 

To evaluate the performance of Uni-Mol+, we tested it on the PCQM4Mv2 dataset from the OGB Large-Scale Challenge [6].

This dataset consists of 370万 training molecules,每一个分子包含三种信息:

  1. each with a SMILES representation
  2. a DFT-generated 3D conformation
  3. and a HOMO-LUMO gap label

输出:小分子

输出: 量子化学性质预测(HOMO-LUMO

本次实验的目的:Uni-Mol+的目的就是通过SMILES预测HOMO-LUMO【由rdkit将2D分子图 --> 3D初始构象 --> 量子化学性质预测

数据集的具体使用具体而言,我们利用PCQM4MV2训练集中提供的DFT平衡构象以及RDKit生成的廉价构象,使用LTI策略来训练Uni-Mol+模型。在PCQM4MV2的验证集中,我们采用RDKit生成的构象作为输入。 

We generated 8 conformations for each molecule using RDKit, at a per-molecule cost of approximately 0.01 seconds. During training, we randomly sampled 1 conformation as input r at each epoch, while during inference, we used the average HOMO-LUMO gap prediction based on 8 conformations.

We used the AdamW optimizer with a learning rate of 2e-4, a batch size of 1024, (β1, β2) set to (0.9, 0.999), and gradient clipping set to 5.0 during training, which lasted for 1.5 million steps, with 150K warmup steps. We also utilized exponential moving average (EMA) with a decay rate of 0.999. The training process required around 5 days, powered by 8 NVIDIA A100 GPUs. Additionally, inference on the 147k test-dev set took approximately 7 minutes, utilizing 8 NVIDIA V100 GPUs.

We incorporate previous submissions to the PCQM4MV2 leaderboard as baselines. In addition to the default 12-layer model, we evaluate the performance of Uni-Mol+ with two variants consisting of 6 and 18 layers, respectively, to investigate the impact of varying model parameter sizes.

数据集具体解释:

wget http://ogb-data.stanford.edu/data/lsc/pcqm4m-v2-train.sdf.tar.gz

这个网址是用于下载OGB-LSC数据集中PCQM4M-LSC任务的训练集原始数据文件。

PCQM4M-LSC(Predicting Chemical Quantum Mechanical (QM) properties using Molecular Machine Learning on Large-scale data Sets)是OGB-LSC任务集合中的一个任务,它是基于图的机器学习领域中的一个经典问题,即预测给定分子的量子力学性质【这个任务需要预测给定分子的热能、配位化学能、电子亲和势等多种量子力学性质】

PCQM4M-LSC数据集是目前规模最大的QM数据集,包含数百万个分子的结构信息和与其相关联的QM性质。使用PCQM4M-LSC数据集进行图神经网络训练和性能评价可以推动分子机器学习领域的进一步发展。

wget https://dgl-data.s3-accelerate.amazonaws.com/dataset/OGB-LSC/pcqm4m-v2.zip

这个网址是下载 OGB (Open Graph Benchmark) Large-Scale Challenge (LSC) 的 pcqm4m-v2 数据集压缩文件。

这个数据集是一个用于图表征学习 (Graph Representation Learning) 的大规模挑战任务,需要参赛选手开发出高效、可扩展、有创意的算法,对一个包含数百万个化学分子的图数据集进行节点表征学习和链接预测。这个数据集主要包含化学分子的结构,以及它们之间的化学键关系和物理性质等信息。

它们之间的区别:

这两个网址都提供了 OGB Large-Scale Challenge (LSC) 的 pcqm4m-v2 数据集,但是下载的数据格式和细节稍有不同。

第一个网址 http://ogb-data.stanford.edu/data/lsc/pcqm4m-v2-train.sdf.tar.gz 只提供了训练集数据压缩包,其中数据的格式是一个包含大量分子信息的 SD 格式压缩包。如果需要验证集和测试集数据,还需要从第二个网址中下载。

第二个网址 https://dgl-data.s3-accelerate.amazonaws.com/dataset/OGB-LSC/pcqm4m-v2.zip 提供了完整的数据集压缩包,其中包含一个总大小为 6.1GB 的 pcqm4m_v2 文件夹,内部包含训练集、验证集和测试集数据、元数据和评估脚本等。其中数据的格式大多是 CSV、JSON 等格式。


二、模型结构 

Uni-Mol+ 是用于量子化学性质预测quantum chemical property predictionQC任务的

在整体框架方面,Uni-Mol+ 基于低成本的方法如RDKit/Openbabel生成初始构象,并通过迭代优化这些构象,使其逼近 DFT 方法得到的高精度稳态构象。这样一来,可以通过基于模型优化后的构象来获得更精确的量子化学性质预测结果。 

Fig 1 . In contrast to prior methods that directly predict QC properties from 1D/2D data, Uni-Mol+ employs a distinct approach.
  1.  Firstly, given a 2D molecular graph, Uni-Mol+ generates an initial 3D conformation from inexpensive methods such as RDKit.【由rdkit将2D分子图生成为3D初始构象】
  2. Then, the initial conformation is iteratively optimized to its equilibrium conformation, and the optimized conformation is further used to predict the QC properties.【将3D初始构象迭代优化为平衡构象,平衡构象用于下一步的QC性质预测】

在模型结构方面,Uni-Mol+ 进一步加强了 Uni-Mol 的双分支 Transformer 结构,以更好地捕捉三维空间的信息。 

Figure 2: The Uni-Mol+ model backbone consists of two tracks of representations - atom and pair, initialized by atom features and 2D graph/3D conformation respectively.

These representations communicate with each other at every block. Besides, Uni-Mol+ optimizes the predicted 3D position iteratively using the previous iteration’s predicted conformations as input for the current iteration.

而在训练策略方面,Uni-Mol+ 提出了一种新的方法,即线性轨迹注入,可以更有效地学习 DFT 构象的优化。

Figure 3: Illustration of the linear trajectory injection (LTI) method for conformational sampling.

The leftmost cheap conformation is generated by RDKit/OpenBabel, while the rightmost conformation is the target DFT conformation. LTI assumes a linear trajectory from the left to the right and enables us to sample conformations by controlling the parameter q. To highlight the differences between the conformations, we include the target DFT conformation as a reference in translucent gray.


三、线性轨迹注入 linear trajectory injection (LTI)

目标:Rdkit得到的廉价构象 --优化为--> DFT平衡构象

方法:为了高效地进行优化构象,我们提出了一种名为线性轨迹注入(LTI)的方法,利用以DFT优化后的平衡构象作为目标。该方法从Rdkit得到的廉价构象和DFT构象之间的路径中选择一个随机构象,并将其作为模型的输入。需要注意的是,在大多数情况下,实际路径是未知的,因此我们使用伪路径(pseudo trajectory)假设两个构象之间的移动是一个线性过程。LTI方法提供了多个分子的输入构象,可以作为有效的数据增强策略。此外,从伪路径中的中间点开始简化了朝向DFT平衡构象的训练。另外,通过迭代优化,LTI模型可以沿着路径逐渐收敛于DFT平衡构象。


值得一提的是,Uni-Mol+仅使用了单个模型,没有使用任何额外的提分技巧。相比之下,之前的冠军方法不仅融合了100多个模型,还额外把验证数据集用于模型训练。从这点来看,Uni-Mol+的模型能力非常强大。

相较于之前的方法,Uni-Mol+ 在整体框架、模型结构和训练策略等方面都进行了创新

  1. 在整体框架方面,Uni-Mol+ 基于低成本的方法如RDKit/Openbabel生成初始构象,并通过迭代优化这些构象,使其逼近 DFT 方法得到的高精度稳态构象。这样一来,可以通过基于模型优化后的构象来获得更精确的量子化学性质预测结果。
  2. 在模型结构方面,Uni-Mol+ 进一步加强了 Uni-Mol 的双分支 Transformer 结构,以更好地捕捉三维空间的信息。
  3. 而在训练策略方面,Uni-Mol+ 提出了一种新的方法,即线性轨迹注入,可以更有效地学习 DFT 构象的优化。


AI4S进入预训练模型时代!Uni-Mol超进化屠榜OGB-LSC,开拓材料设计新天地