具有分子结构多重图的分子动力学驱动图神经网络

作者 | 刘帅

审核 | 宋从智

今天介绍纽约城市大学团队发表在NeurIPS 2020上的文章“Molecular Mechanics-Driven Graph Neural Network with Multiplex Graph for Molecular Structures”。在这项研究中，作者针对分子性质预测问题，引入了分子动力学驱动的理念，设计了双层多重图来对分子建模，同时捕捉分子的局部和全局信息。通过信息传递模块，模型的表达能力和计算复杂度得到了很好的平衡。实验证明作者提出的方法优于各基准方法，并且可以学习到更丰富的分子表示。

1.摘要

从分子结构预测理化性质是人工智能辅助分子设计的关键任务，目前已经提出了越来越多的图神经网络 (GNN) 来应对这一挑战。这些模型通过在分子中加入辅助信息来提高它们的表达能力，同时不可避免地增加了它们的计算复杂度。在这项工作中，作者的目标是设计一个对分子结构既强大又高效的 GNN。为了实现这一目标，作者提出了一种分子动力学驱动的方法，首先将每个分子表示为一个两层多重图，其中一层仅包含主要捕获共价相互作用的局部连接，另一层包含可以模拟的全局连接非共价相互作用。然后对于每一层，提出相应的消息传递模块来平衡表达能力和计算复杂度。基于这两个模块，作者构建了多重分子图神经网络 (MXMNet)。通过在 QM9 小分子数据集和大蛋白质配体复合物的 PDBBind 数据集验证时，MXMNet 在资源受限的情况下取得了优于现有最先进模型的结果。

2.模型框架

作者首先对分子动力学方法从几何的角度进行分析，发现对于传统局部和非局部关系模型，可以对应转化为对局部和非局部几何信息，其中局部几何信息包含局部角度和近邻距离，而非局部几何信息包含近邻距离和远邻距离。由此，对于一个分子图，可以构建出一个两层多重图——一层 (Local GI) 代表着局部几何信息，另一层 (Global GI) 代表着全局的几何信息。同时用不同的层次对其进行消息传递，这样可以更好兼顾两方面的信息，同时降低了计算复杂度。

具体的层次设计如上图所示，(a)为作者用来学习两层多重图的MXM模块，(g)为堆叠MXM模块形成的MXMNet，而(b)(c)(d)(e)(f)分别为模块内部的具体设计。

首先，针对全局几何信息的传递，作者主要通过两节点之间的边设计传递方式（如图(b))：

同时，在两次信息传递之间，作者引入了多个残差块（如图(d)(e))，来增强表示的表达能力。

其次，针对局部几何信息传递（2-hop），作者结合了角度和两节点间距离进行设计。如上图所示，边可以对应两种角度——1-hop边和2-hop边之间的夹角为2-hop角，而1-hop边之间的夹角为1-hop角。因此，作者提出了三步进行局部几何信息传递的方式（如图(c))：

简单来看，作者先捕捉2-hop角和对应边及节点信息，再捕捉1-hop角和对应边及节点信息，最终进行信息更新，并进行预测输出。

为了更好地结合两个层次的信息，作者引入了交叉层 (Cross Layer)，如下:

交叉层以上一层的节点表示为输入，将其转换成下一层的输入，形成图(a)中交叉学习的方式。

最终，作者将多个MXM模块叠加在一起，同时将输入的分子信息embedding化，包括每个原子的表示和节点间距离和角度的表示（RBF & SBF模块）。最终将每个MXM模块的输出统一，形成最后的输出。

4.实验结果

作者主要使用了QM9和PDBBind分子数据集，在分子性质预测问题上展开模型性能对比实验。在QM9数据集的具体性能对比如下。可以发现，作者提出的方法表现强劲，在不同的目标上均有不错的表现。同时，作者对batchsize大小和全局距离值设置进行了探究，可以看出，batchsize的影响不大，全局距离值应该根据不同任务进行不同的设置。

作者进一步进行了消减分析。不难发现，两种传递方式缺一不可，同时使用它们才能得到最好的性能。

在PDBBind数据集的具体性能对比如下。可以看出，MXMNet优于所有Baseline，效果出众。

5.总结

在本文中，作者提出了一个两层多重图来对分子的局部和非局部几何信息进行建模。针对两种模式的不同特点，设计了两种信息传递方式，有效地捕捉到了关键信息，同时降低了计算复杂度。最终实验证明，MXMNet优于现有方法，同时计算效率高，表示能力强。