3D Infomax improves GNNs for Molecular Property Prediction

出处

作者：Hannes Stark等
机构：Massachusetts Institute of Technology等
期刊：Proceedings of the 39 th International Conference on Machine Learning，2022/06/04
Code ：github

摘要

使用现有的三维分子数据集来预先训练一个模型，以推理出仅有二维分子图的分子的几何形状。
模型的名称为 3D Infomax，最大化学习到的3D summary vector和GNN的表征之间的相互信息(mutual info)。
使用未知几何形状的分子进行微调，GNN仍然能够提供一些隐性的3D信息并用于下游任务。
在很多属性上有着较大的进步，比如在QM9量子力学特性上，MAE减少了22%

介绍

现有的分子特性预测方法和3D infomax的动机：

标准方法：利用GNN和2D的分子图，结果快但差；
显性的3D方法：使用经典的方法或者机器学习的方法计算3D坐标，然后作为输入进行预测。结果准确但是对于实际应用来说计算坐标太慢。
3D Infomax：① 预训练：用一个2D网络对有3D信息的分子进行训练，得到有着隐性3D信息的表征。② 将2D网络的参数微调。结果真是又快又好。

背景

2D分子图
- $G=(ν,ϵ)，其中ν是节点−原子，ϵ是边−共价键G=(\nu,\epsilon)，其中\nu是节点-原子，\epsilon 是边-共价键$ ，边可以包含键类型信息，节点可以包含一些特征数据，比如原子编号，但至此都无3D坐标信息。
3D 分子构象
- 不同构象会带来不同化学性质，为保证抓住3D信息，所以需要考虑几乎所有的构象。
- 当考虑一个分子 $c$ 个已知的构象时，把他们表示成一组点云 ${R}j1⋯cj\{R\}^j_{j1\cdots c}$ 。每个点云 $R={rv}v∈νR=\{r_v\}_{v \in \nu}$ 表示分子中所有原子 $ν\nu$ 的坐标（即一组点云是一个构象的所有坐标点集合）。
- RD-Kit的ETKDG算法能快速计算构象但是不准确；最流行的是CREST，速度和准确率兼备，但仍然需要大约6小时（per cpu）完成一个药类分子的计算。
分子的对称性
- 当所有原子坐标 jointly translated或者围着一点旋转（SE(3)对称），那么分子的构象就不会改变。同时，分子的性质会被他们的手性决定。我们的方法也能在表征中体现对称性。
图神经网络
- 大部分GNN可以被一个MPNN框架描述，比如我们用的PNA模型。
- MPNN的目的是为了学习一个图的表征。他们通过不断迭代地去应用消息传递层，然后将所有点的表征结合。一个消息传递层通过使用置换不变性函数（mean，max，sum，不管数据如何置换位置，结果都不变）计算该点的邻点和其之间的边的值以用于更新该点信息。消息传递层之后，另一个置换不变性函数被用于提取点层的embedding到图层的embedding。

方法（3D Infomax）

使用对比学习完成了输入为2D信息，但可以推断出3D几何信息的模型。预训练模型为PNA（将图多种方法聚合，sota和简单）

图1

最大化一个使用2D分子图的2D GNN和使用3D构象的3D GNN之间的互信息（在机器学习中，理想情况下，当互信息最大，可以认为从数据集中拟合出来的随机变量的概率分布与真实分布相同。）
与训练之后，我们将它迁移到属性预测任务，并微调。在微调过程中，GNN产出的3D信息会被用于提高预测。
在图1中，有两个模型。需要预训练的是2D网络 $network f^a$ ，它可以产出一个表征 $fa(G)=za∈Rdzf^a(G)=z^a \in \R^{d_z}$ ；另一个将 $R={rv}v∈νR=\{r_v\}_{v \in \nu}$ 编码的3D网络 $network f^b$ 给出一个表征 $fb(G)=zb∈Rdzf^b(G)=z^b \in \R^{d_z}$ 。可以当成是一个对比蒸馏，因为student 2D网络可以从teacher 3D网络那学会生产3D信息。
对比框架
为了教会2D $network f^a$ 从2D图输入中学到3D信息，我们最大化了潜在2D表征 $z^a$ 和3D表征 $z^b$ 的互信息。因为当两者来自同一个分子，那我们希望 $z^a$ 和 $z^b$ 尽可能的一致，所以利用了图2的对比学习。
对于一组batch，中间包含N个分子图 ${Gi}i∈{1⋯N}\{G_i\}_{i \in \{1\cdots N\}}$ ，点坐标 ${Ri}i∈{1⋯N}\{R_i\}_{i \in \{1\cdots N\}}$ ，然后得到多个表征 $z_i^a$ 和 $z_i^b$ 。
图2
对比学习的第一个目的就是如果两者是正样本对，那么就要最大化表征的相似度，表示他们是来自同一个分子（同一个index i）。第二个目标就是强迫负样本对 $z_i^a$ 和 $zkb,i≠kz_k^b, i\ne k$ 不相似。
这两个目标都是通过修改NTXent loss实现（如何实现相似的越相似）的：

其中

是余弦相似度， $τ\tau$ 是一个温度参数（超参），可以当做最相似的负样本对的权重（也就是当负样本对很相似时，调整 $τ\tau$ ）。不同的对比损失组合和自监督学习是有可能学会一个2D和3D表征之间的联合嵌入空间，上面的函数是表现最好的。
使用多构象
使用 $c$ 个最高概率的第 i 个构象 ${Rij}j∈{1⋯c}\{R_i^j\}_{j \in \{1\cdots c\}}$ ，如果不够c个就把能量最低的重复。图2右边，就是将分子的2D表征和每一个构象进行比较。

3D网络

3D网络输入是原子坐标作为点云，然后输出一个置换不变的表征 $z^b$ ，尽可能多的把3D结构信息编码，但是不能够接触2D信息，不然的话互信息可能会因为两个模型的交互变得更大。
我们的模型将每一对atom的欧几里得距离进行编码，这样表征可以定义所有原子的相对位置并且保证旋转平移不变性，并且也是反射不变的，但是对手性分子没办法区别。
$d_{uv}$ 表示u原子和v原子之间的距离，会先使用高频的sine和cosine去投影到一个高维空间（因为键长之间的区别比较小）。然后以F=4的频率map（有点类似position encoding），更详细的操作可见MPNN框架。

数据

3D数据集是来自QM9（134k个平均18个原子的只有一个构象的小分子，kaggle下载）、GEOM-Drugs（304k）和QMugs（665k）。后两者有较大的且是多构象的药类分子（44.4和30.6平均原子个数）
微调：预测十个来自QM9和GEOM-Drug的量子特性，这些数据不与预训练的数据相交。
预训练用了50k单构象来自QM9，140k5构象来自GEON-Drugs，620k3构象来自QMugs

对比

Baseline

距离预测：使用已有的最低能量构象去预训练一个GNN，以直接预测所有原子之间的距离。然后将任意两个u,v原子的表征简单地拼接在一起（uv,vu），随后放入mlp（U，直接降到1维），||表示拼接

$softplus(x) = log(1 + e^x )$ ，loss function是MSE。
构象生成：GeoMol（sota生成分子构象的模型），一个生成式模型，产生一个分子的可能的3D结构的分布，从而获取到多构象信息。利用他们模型做预训练任务然后提取网络用于不同下游任务
GraphCL：一个卷积增强预训练模型with JOAO配置，模型通过学习产出一个对增强不变的表征来完成自监督目标。
结果
数值为MAE，RAND INIT模型随机初始参数，PROPRED指用GEOM-Drugs的Gibbs自由能来做的预训练，DISPRED指用有最高概率的构象去预测所有原子的距离，CONFGEN指与训练的时候预测10个构象，3D Infomax分别使用三个数据集做预训练，RDKIT SMP使用RDKit生成的3D坐标输入SMP（一个GNN）做训练，True 3D SMP最后一列是用真实的3d坐标使用SMP预测的，蓝色表示improvement，橙色表示worse。

对QM9数据集中的8中特性做预测

结论

相当于一个2D分子图的预训练模型，能得到隐含3D信息的表征，并且具有一定泛化能力（不会有负迁移），可以借助同一个分子的多构象信息来帮助下游属性预测任务。

【读文献】3D Infomax 小分子预训练模型相关推荐

【综述】分子预训练模型综述
A Systematic Survey of Molecular Pre-trained Models 目录总结一.Introduction 二.Molecular Descriptors 三.P ...
一文读懂最强中文NLP预训练模型ERNIE
基于飞桨开源的持续学习的语义理解框架ERNIE 2.0,及基于此框架的ERNIE 2.0预训练模型,在共计16个中英文任务上超越了BERT和XLNet, 取得了SOTA效果.本文带你进一步深入了解ER ...
腾讯优图开源业界首个3D医疗影像大数据预训练模型
整理 | Jane出品 | AI科技大本营(ID:rgznai100) 近日,腾讯优图首个医疗AI深度学习预训练模型 MedicalNet 正式对外开源.这也是全球第一个提供多种 3D 医疗影像专用预 ...
全球首个！腾讯优图开源3D医疗影像大数据预训练模型
点击我爱计算机视觉标星,更快获取CVML新技术本文转载自腾讯优图. 近日,腾讯优图首个医疗AI深度学习预训练模型MedicalNet正式对外开源.这也是全球第一个提供多种3D医疗影像专用预训练模型的 ...
香侬读 | 让预训练模型学习知识：使用多学习器增强知识建模能力
论文标题: K-Adapter: Infusing Knowledge into Pre-Trained Models with Adapters 论文作者: Ruize Wang, Duyu Tan ...
让预训练模型学习知识：使用多学习器增强知识建模能力
论文标题: K-Adapter: Infusing Knowledge into Pre-Trained Models with Adapters 论文作者: Ruize Wang, Duyu Tan ...
翟季冬：基于国产超算的百万亿参数超大预训练模型训练方法
[前沿进展]训练参数规模万亿的预训练模型,对于超级计算机而言是不小的挑战.如何提升超算的计算效率,实现更大规模的参数训练,成为近年来研究者探索的课题.在近日举办的Big Model Meetup第二期 ...
NLP千亿预训练模型的“第四范式”之Prompt Learning介绍分享
点击上方"AI遇见机器学习",选择"星标"公众号重磅干货,第一时间送达论文转载自知乎专栏:ai炼丹师作者:避暑山庄梁朝伟一.背景随着GPT-3诞生,最 ...
腾讯开源首个医疗AI项目，业内首个3D医疗影像大数据预训练模型
乾明发自凹非寺量子位报道 | 公众号 QbitAI 腾讯AI,开源又有新动作. 旗下顶级AI实验室腾讯优图,对外开源了腾讯首个医疗AI项目--深度学习预训练模型MedicalNet. 这一项 ...

【读文献】3D Infomax 小分子预训练模型

3D Infomax improves GNNs for Molecular Property Prediction

出处

摘要

介绍

背景

相关文献

方法（3D Infomax）

3D网络

数据

对比

结论

【读文献】3D Infomax 小分子预训练模型相关推荐

最新文章

热门文章