在人工智能领域，深度学习平台的重要性毋庸置疑。飞桨是百度自研的深度学习平台，飞桨社区的开发者基于飞桨平台积极的将自己的智慧应用到前沿技术的研究中，为人工智能长期研发目标贡献自己的力量。为此，飞桨开发者说专栏推出顶会专题，给予飞桨社区开发者一个平台，介绍学术研究的最新发展动向，分享在学术研究方面的经验。
本期文章将为大家分享飞桨社区开发者李双利、周景博发表于AAAI2022的文章《GeomGCL: Geometric Graph Contrastive Learning for Molecular Property Prediction》
飞桨社区核心开发者，百度研究院商业智能实验室研究实习生，中国科学技术大学的在读博士生。
主要进行时空数据挖掘和图深度学习的相关研究工作。曾获2021年百度研究院年度优秀实习生，有多项基于飞桨完成的论文，发表于KDD、AAAI等计算机顶级会议。
飞桨高级开发者技术专家（高级PPDE），现任百度研究院商业智能实验室资深研究员。
主要从事数据挖掘和机器学习相关的研究和应用工作，包括时空大数据、深度几何学习、知识图谱和AI辅助药物设计等，PaddleSpatial技术负责人，基于飞桨完成论文多篇，发表于KDD、AAAI、TKDE等计算机顶级会议和期刊上。
在计算生物和计算化学领域，准确预测分子的各种生化性质的准确预测对于药物研发等应用领域至关重要。然而在真实场景下，大部分数据往往是缺少标签的，传统的有监督学习算法无法利用这些数量巨大的缺失标签的分子进行学习，因此如何用图对比学习的自监督策略增强模型的表达能力是十分有必要的。
然而，当前自监督的分子性质预测方法基本都是基于二维分子图设计的预训练任务或者进行分子图的增强的方法，一方面有效的预训练任务往往需要生物化学领域的专家知识进行精心设计，而分子图增强策略基本都会破坏分子本身的结构性质，产生的分子很可能是没有化学意义的，不利于模型的对比学习。另一方面，分子本质上是一个由多种类型原子相互作用构成的网络结构图，与一般的研究对象不同，分子图除了拓扑结构信息，其中还包含关键的空间结构信息。当前大多数无监督的分子性质预测方法只考虑了二维视角的分子信息，没有将三维空间的分子图加入自监督的学习框架进行有效利用。而原始的分子化学式可以转化为二维平面结构图和三维空间结构图，如何充分学习结合两种视角下的信息对于分子的表征学习和性质预测是十分有益的。
图1 二维视图和三维视图中的空间几何信息
在本文中，我们提出了一个全新的分子图对比学习框架，可以对大量无标签的分子数据进行自监督的学习，同时考虑了二维和三维结构下的分子图的拓扑信息和空间结构信息，基于图神经网络有效融合不同视角的空间特性对分子表征进行有效学习。
本方法是基于飞桨2.3版本实现模型训练。
首先，使用飞桨的PGL图学习框架进行二维和三维视图的高效构建，通过PGL提供的异构图message passing机制可以便捷的进行分子图相关的空间结构学习，在第一阶段首先使用空间对比学习进行飞桨框架上的预训练。
第二阶段基于预训练模型根据不同的分子性质预测任务进行微调，显著提升了分子性质的预测效果。
图2 二维-三维分子图对比学习的模型框架
如图2所示，给定一个SMILES形式的分子化学式，我们首先转换得到二维和三维两种视图的分子图。针对空间属性的表征学习，我们通过飞桨定义使用径向基函数（RBF，radial basis functions）来对空间几何信息进行编码得到二维和三维的几何向量表征，代码实现如下所示。
如下代码所示，我们基于空间信息表征向量提出了自适应的几何消息传递框架GeomMPNN来通过“节点—边”（Node-Edge）的交互方式实现同时高效学习分子的拓扑结构和空间几何信息。
基于Paddle Graph Learning (PGL) 框架，我们可以实现多层分子图交互过程。PGL是一个基于飞桨的高效易用的图学习框架，借助于简便的消息传递范式可以作为模型进行多级原子节点和边表征学习的接口
。总体来看，模型包含三个消息传递层：Node⟶Edge, Edge⟶Edge和 Edge⟶Node。这三层交替更新节点的表征和边的表征，叠加多层最后使用Node⟶Graph的注意力池化层来得到最终的分子表征。
由于只有化学键对应的边具备初始特征，我们首先对每一对节点（原子）的表征进行聚合来更新得到边的表征。同时为了从二维和三维两个方面增强联系性，对聚合信息进一步分别融合共价键特征和三维距离表征。
在得到边表征后，我们进一步设计了边到边的消息传递过程来捕捉原子节点的空间分布。空间角度感知的聚合策略可以分别作用于二维和三维关系边，基于PGL的send-recv消息传递实现机制可以方便的实现二维边之间的消息传递函数。
对于三维视图而言，其边的分布相对稠密，因此三维边周围存在较多的边邻居，我们基于三维空间角度 θ 把每个目标三维边的邻居边切分到不同角度邻域 A_1,…,A_n，接着我们采用局部到全局的层级聚合过程来学习三维视图下的边表征：
这种全局策略可以从更好的提取三维空间分布信息来加强分子的几何结构信息的学习。
经过上一阶段得到了角度信息感知的二维和三维边表征〖 e〗_uv^(2d,t) 和 e_uv^(3d,t) 后，我们进一步采用了Edge⟶Node 消息传递层来完成从边到节点的信息传回过程。该过程设计了类似于上一阶段的自适应交互过程来学习节点之间的距离几何特性：
由于以上交互式的消息传递过程可以对二维和三维视图中角度和距离这些几何特性进行有效学习，在叠加T层之后得到的节点（原子）表征 a_v^(2d,T)和 a_v^(3d,T)充分包含了分子图的拓扑结构信息和空间几何信息。为了进一步得到分子表征，我们使用PGL的图计算方法进行基于注意力机制的池化过程来识别重要的原子节点。
最后构建对比学习的正负样本，对于每一个批次的输入，我们把同一分子生成的二维和三维视图作为一对正样本 {z_i^2d,z_i^3d }，而其他分子的视图表征则都作为负样本。模型的学习目标是最大化正样本之间的相对一致性，使负样本之间尽可能差距变大从而实现模型训练。
我们在多个分子数据集上的实验证明了所提出的对比学习框架可以显著提升分子性质的预测效果，表1表明了融合二维和三维分子视图不仅比有监督方法效果更佳，而且预测效果优于最新的无监督预训练模型和图对比学习方法。
此外，我们通过对二维和三维视图学习模块进行消融实验证明了充分融合两个维度的分子信息才能达到模型最优的预测效果，并且利用对比学习策略可以进一步提升效果。
图3 消融实验
本文提出了一种新的二维-三维分子图对比学习框架，通过设计的双通道几何信息传递神经网络来充分捕获二维和三维视图下的距离和角度信息，然后提出了分子几何对比学习策略提升分子表征学习能力，多个分子数据集的实验结果证明了所提出方法的有效性。
*相关代码已经开源在飞桨螺旋桨（PaddleHelix）生物计算平台上。螺旋桨PaddleHelix生物计算平台是基于百度飞桨深度学习框架开发的生物计算平台，提供AI+生物计算能力，满足新药研发、疫苗设计、精准医疗场景的AI需求。
https://ojs.aaai.org/index.php/AAAI/article/view/20377
https://github.com/PaddlePaddle/PaddleHelix/tree/dev/research/geomgcl
关注【飞桨PaddlePaddle】公众号
获取更多技术内容~
