星药科技发表顶会论文 | CoMPT：基于Transformer型图交互式神经网络的属性分子图表征学习方法

2021-08-13 16:04

可视化研究也进一步表明该模型实现了更好的药物分子图表征能力。

来源 | 星药科技

近日，中山大学与星药科技在国际人工智能顶级会议 IJCAI-2021 上联合发表了论文：Learning Attributed Graph Representations with Communicative Message Passing Transformer。

本文提出了CoMPT（Communicative Message Passing Transformer）神经网络，一种通过在Transformer架构中融合消息传递机制并加强节点与边之间的消息交互来改进药物分子图的表征方法。此外，受热扩散现象的启发，该表征方法通过将消息传递机制转化为消息扩散机制减少了消息过度富集的影响。

大量实验表明，本文提出的CoMPT模型在七个化学性质数据集（图级任务）和两个化学位移数据集（节点级任务）上的性能优于基线模型的性能（平均性能提升约4%）。可视化研究也进一步表明该模型实现了更好的药物分子图表征能力。

研究背景

如何构建合适的分子图表征是材料科学、化学和药物设计等众多领域的核心问题之一。近年来研究者尝试将分子抽象成带有标注属性的结构图，再利用图神经网络(GNN)进行分子图的表征学习，取得了比较显著的成果。

1.1

存在的问题

尽管当前已有的模型表征功能强大，但它们要么基于图中节点的局部聚合操作而丢失图的高阶属性，要么只关注节点信息而没有充分利用边信息。这些问题都使得GNN在分子图的表征学习任务中具有极大的局限性。目前已经有一部分工作尝试去解决以上问题，分为两种不同的思路：一种是引入Transformer型的GNN来学习分子图结构数据中的远程依赖，但它忽略了分子图中局部子结构的归纳偏置，因而在一些以拓扑结构为图性质的重要影响因素的任务中表现不佳；另一种是通过引入边的方向性改进原有的消息传递神经网络来提升性能，这种方法虽然利用了边信息，但仍然无法处理远程依赖。

1.2

作者的解决思路

作者提出了一种用于分子表征学习的Transformer型图交互式神经网络(CoMPT)。与之前只着重于节点信息的Transformer型GNNs相比，CoMPT在Transformer架构中融入节点和边的嵌入消息交互过程，然后该过程所产生的新消息反过来用于更新节点和边的嵌入。此外，通过结合拓扑连接矩阵来引入消息扩散机制改进原来的消息传递机制，从而减少节点和边的消息富集冗余。以上这些机制使得模型可以在分子图中选择性地传播消息，CoMPT能够为下游预测任务提取更准确且富有意义的分子图表征。这项工作的主要贡献总结如下：

提出了一种新的用于分子表征学习的Transformer型图交互式神经网络，即CoMPT，它能够对分子图的原子和键进行消息交互和消息更新，从而可以同时捕获分子图的局部和全局结构信息。

受热扩散现象的启发，该模型利用拓扑连接矩阵将原来的消息传递机制转化成消息扩散机制，这种机制被进一步证明可以缓解GNN中普遍存在的过度平滑问题。

在图级和节点级数据集上进行的实验证明了该模型的有效性。CoMPT在所有九项任务中的表现都超越了之前性能最优异的模型，平均性能提高了约4%。

模型概述

2.1

问题定义

分子图结构可以表示为

G=(\mathcal{V},\mathcal{E})

，其中

\left|\mathcal{V}\right|=n

表示这个分子有

n

个原子(节点)，

\left|\mathcal{E}\right|=m

表示有根化学键(边)。

\mathcal{N}_v

表示节点

v

的邻居节点集合。对于每个节点和每条边，分别使用

f_{node}

和

f_{edge}

表示其初始特征向量的维度。我们沿用DMPNN中使用的有向分子图模型以减少不必要的循环信息传递。我们使用

x_v\in\mathbb{R}^{f_{node}}

表示节点的初始特征向量，使用

e_{uv}\in\mathbb{R}^{f_{edge}}

表示有向边

u\rightarrow v

的初始特征向量，大写字母

X\in\mathbb{R}^{n\times f_{node}}

以及

E\in\mathbb{R}^{n\times{n\times f}_{edge}}

分别表示节点特征矩阵和边特征矩阵。除此之外，分子图表征学习任务通常分为如下2大类4小类问题：(1) 图分类/回归问题，这类问题中给定一系列的分子图{

{G_1,\ldots,G_N}

}以及对应的标签 {

{y_1,\ldots,y_N}

}，任务是预测一个新的分子图对应的标签。(2) 节点分类/回归问题，这类问题中给定分子图中多个节点的标签，任务是预测新分子图中节点的标签。

2.2

CoMPT模型

CoMPT模型的核心算法是使用节点-边消息交互模块取代Transformer框架中的自注意力模块来完成消息传递机制并学习分子图的嵌入表征。

图1.MPNN, Transformer与CoMPT的消息传递流程对比

CoMPT模型需要三个输入：节点初始特征矩阵

X

、边初始特征矩阵

E

，以及使用任意两个节点之间的最短路径长度计算得到的拓扑连接矩阵

A

。受原来的Transformer框架和有向消息传递机制的启发，研究人员在原始输入特征中为原子特征添加位置信息，为边特征添加方向信息。具体来说，对于每一个原子

v_i(i=1,2,\ldots,n)

，我们根据原子索引学习一个权重向量，公式如下：

其中

node_embedding(\cdot)

将原始特征向量的维度映射到给定维度以保持一致。对于每一条边

e_{uv}

，我们使用课题组以前发表的CMPNN模型中使用的方法，即将起始节点特征向量加到边特征向量中，公式如下：

其中

edge_embedding(\cdot)

将原始边特征向量映射到给定维度以保持一致。为了描述方便，我们设定这个统一的隐藏层维度为

f

。

CoMPT的核心算法思想是在Transformer框架中使用节点隐藏特征矩阵

h(X)

和边隐藏特征矩阵

h(E)

来计算消息交互分数矩阵

M

，它取代了原始Transformer Encoder层中的自注意力分数矩阵。具体来说，首先根据

h(X)

计算矩阵，根据

h(E)

计算矩阵

K

，公式如下

然后为了适配点边交互矩阵的维度，我们使用了爱因斯坦求和约定符号计算如下三个消息交互分数矩阵

M

从计算原理角度来说，爱因斯坦求和约定符号中的第一步与原始Transformer框架中计算

QK^T

的过程完全相同，在这里即为点边交互全消息矩阵

T

，它的每个位置元素表示任意一个节点与任意一条边交互得到的信息。

接下来，该符号的第二步根据输入参数要求取相应位置的消息，在这里我们选择了每个节点与其出向边交互的消息矩阵

M_o

，与其入向边交互的消息矩阵

M_i

，与其自环边交互的消息矩阵

M_d

，三个部分的信息后再利用softmax函数归一化，得到最终的消息交互矩阵

M

该操作消除了重复计数的自循环消息，避免了信息冗余。该消息交互矩阵

M

进一步用于更新每个Encoder层中的节点隐藏特征矩阵

h\left(X\right)

和边缘隐藏特征

h(E)

这里的

\odot

表示元素乘。除此之外，CoMPT沿用原始Transformer框架中的“多头”消息交互机制。在堆叠多个CoMPT的Encoder层的时候，我们为了使得训练过程更加稳定，使用了前置层归一化技术(post layer norm)。另外，在连接相邻两个Encoder层的时候，为了减少梯度消失的问题，我们还使用了残差技术：

这里的

k

表示层号索引。

2.3

消息扩散机制和全局池化

在消息传递机制中，保持消息交互的正确性是准确预测图级/节点级任务的属性的关键，但以前的研究表明在GNN中过度深层的消息聚合会导致过平滑现象(分子图中所有的原子嵌入表征趋近于相同的值的现象)。为了减缓这种现象发生的程度，研究人员受热扩散现象的过程启发，为消息传递机制的迭代过程设计了一个简单的消息衰减机制，以减缓过度冗余的信息聚合到节点和边的过程。前面提到的消息交互得分矩阵

M

中的每个隐藏向量都可以看作是行索引指代的节点准备传给列索引指代的节点的消息，而拓扑连接矩阵A中对应位置的元素则表示两个节点之间在分子图中的最短拓扑路径距离。因此我们可以使用指数衰减因子来减少待传递的消息：

M\left(u,v\right)=M\left(u,v\right)e^{-\alpha A(u,v)},

其中

u,v

表示两个节点，

M(u,v)

表示节点

u

要传向节点

v

的消息，

A(u,v)

表示在分子图中节点

u

到节点

v

的最短拓扑路径距离，

\alpha\in[0,1

表示可学习的衰减系数。从公式中就可以看到，随着距离的增加，待传递的消息在开始时会迅速衰减，然后在很长的一段距离范围内都变得平稳。因此，应用这种机制就可以在一定程度上延缓消息传递过程中发生的过度平滑现象。

实验方法和结果

研究人员在三类图数据集(图分类任务，图回归任务，节点回归任务)上对CoMPT模型进行了评估，以回答下面3个问题。

RQ1：与此前最先进的分子图性质预测方法相比，CoMPT 模型性能表现如何？

RQ2：不同的子模块（即节点位置模块、边方向信息模块和消息扩散机制）如何影响 CoMPT模型的预测性能？

RQ3：CoMPT模型能否为分子图结构提供更好的表征形式？

3.1

基准数据集

研究人员共使用了9个基准数据集，包括用于图分类任务的BBBP、Tox21、Sider和ClinTox，用于图回归任务的ESOL、FreeSolv和Lipophilicity，以及用于节点回归任务的分子碳氢化学位移预测。

3.2

模型对比

在图级任务中，研究人员将CoMPT模型与12 种基线方法进行了全面比较。这些模型包括：TF_Robust（一个基于 DNN 的多任务模型）；GCN、Weave 和 SchNet（三个基于图卷积的模型）；N-Gram（一种用于分子性质预测的无监督学习模型）；AttentiveFP（一种基于图注意力网络的模型）；MPNN 及其变体MGCN、DMPNN和CMPNN（四个基于消息传递框架及其变体扩展的模型）；Smiles Transformer和GROVER（两个基于transformer框架的模型）。

在节点级任务中，将CoMPT与HOSE、GCN和MPNN这三种方法进行了比较。

3.3

实验结果

3.3.1 模型性能比较（RQ1）

在图级任务中的性能

据表1结果显示：在7个图级任务的数据集中，CoMPT在其中6个都取得了最佳性能，表明CoMPT模型能够有效地学习到分子图的结构表征并进行图级水平的性质预测。另外还能观察到以下3个现象：(1) 基于消息传递机制的模型和基于Transformer框架的模型的性能表现均优于基于图神经网络的模型，而CoMPT模型融合了这两者的优点，因此性能表现更加优异。(2) 基于消息传递机制的模型普遍比基于Transformer框架的模型性能表现更好，这表明在分子图的结构学习当中，消息传递机制使用的边特征矩阵远比Transformer框架使用的邻接矩阵或者距离矩阵重要，从侧面反映出边特征的重要性。(3) 在如FreeSolv等小数据集上CoMPT 模型依旧可以取得较优异的性能表现，这表明CoMPT利用很少量的数据就可以提升预测性能。

表1.CoMPT和其他模型在七个图数据集上的预测性能结果。灰色单元格表示以前的最佳方法，粗体样式的单元格表示CoMPT获得的最佳结果

在节点级任务中的性能

表2结果表明：CoMPT模型可以提取分子图中有意义的节点嵌入表征，从而能够更准确地预测新分子的NMR谱。

表2.节点级任务的性能

3.3.2 模型消融研究（RQ2）

研究人员对三个基准数据集进行了消融研究以探索影响CoMPT模型性能的因素。如表3所示：集成了节点位置模块，边方向信息模块以及消息扩散机制模块的CoMPT模型在所有架构中表现出最佳性能。其中去除掉消息扩散机制模块导致的性能下降最为明显，这表明了在CoMPT的消息传递过程中减少节点和边消息富集冗余的重要性。除此之外，节点位置模块以及边方向信息模块的使用都有助于提升最终的预测性能。

表3.三个数据集上的消融结果

3.3.3 原子表示可视化（RQ3）

研究人员使用t-分布领域嵌入算法 (t-SNE) 对3个具有共同主要骨架（四元环）、不同侧链的相似化合物的原子嵌入分布进行了可视化展示（图2）。在理想情况下，这三个分子的主要骨架部分的原子嵌入分布应该一致，而独特侧链的原子嵌入分布应该不同。可视化实验结果表明，CMPNN由于受到了过度平滑性的限制，每个分子的所有原子都趋向于一致，CoMPT模型能够很好地分散每个分子不同原子的嵌入，其中带有消息扩散机制的CoMPT模型结果符合预期，即主要骨架部分的原子嵌入分布一致，侧链部分的原子嵌入有所不同。另外，比较有趣的是CoMPT能够对同样的官能团给出不同的原子嵌入分布，如图中第2个分子和第3个分子中对于羟基酮基团的原子嵌入分布(黄圈和绿圈)。这些观察表明了CoMPT模型不仅可以缓解过度平滑现象，而且可以在分子图的更广泛背景下捕获更好的表征，且该模型在节点级任务中也展现了更好的预测性能。

图2.三个相似分子（三种颜色）的原子嵌入的T-SNE可视化

本文提出了一种基于Transformer型图交互式神经网络(CoMPT)，通过在Transformer架构中融合消息传递机制并加强节点与边之间的消息交互来改进药物分子图的表征方法。此外还引入了消息扩散机制来改进消息传递机制，以衰减消息传递过程中节点和边的信息富集冗余，并解决由此带来的过度平滑问题。大量的实验表明，本文提出的CoMPT模型在图级任务和节点级任务上均获得了相对于基线模型更为卓越的性能表现。

参考资料：

Chen, J., Zheng, S., Song, Y., Rao, J., & Yang, Y. (2021). Learning Attributed Graph Representations with Communicative Message Passing Transformer. arXiv preprint arXiv:2107.08773.

原文链接：https://arxiv.org/abs/2107.08773

代码链接：https://github.com/jcchan23/CoMPT