[论文解读] Learning Attributed Graph Representations with Communicative Message Passing Transformer

Jianwen Chen, Shuangjia Zheng+3 more|arXiv (Cornell University)|Jul 19, 2021

Computational Drug Discovery Methods参考文献 17被引用 10

一句话总结

本文提出 CoMPT，一种基于 Transformer 架构的通信式消息传递图神经网络，通过显式建模节点与边之间的双向消息传递，增强分子图表示。通过引入拓扑消息扩散机制，缓解消息爆炸与过度平滑问题，CoMPT 在图级别任务中平均性能提升约 4%，并在 1H-NMR 节点级别预测中达到 0.214 MAE 的 SOTA 表现。

ABSTRACT

Constructing appropriate representations of molecules lies at the core of numerous tasks such as material science, chemistry and drug designs. Recent researches abstract molecules as attributed graphs and employ graph neural networks (GNN) for molecular representation learning, which have made remarkable achievements in molecular graph modeling. Albeit powerful, current models either are based on local aggregation operations and thus miss higher-order graph properties or focus on only node information without fully using the edge information. For this sake, we propose a Communicative Message Passing Transformer (CoMPT) neural network to improve the molecular graph representation by reinforcing message interactions between nodes and edges based on the Transformer architecture. Unlike the previous transformer-style GNNs that treat molecules as fully connected graphs, we introduce a message diffusion mechanism to leverage the graph connectivity inductive bias and reduce the message enrichment explosion. Extensive experiments demonstrated that the proposed model obtained superior performances (around 4$\%$ on average) against state-of-the-art baselines on seven chemical property datasets (graph-level tasks) and two chemical shift datasets (node-level tasks). Further visualization studies also indicated a better representation capacity achieved by our model.

研究动机与目标

• 解决现有 GNN 在捕捉长距离依赖关系以及在分子图中未充分利用边信息方面的局限性。
• 克服深度 GNN 和基于 Transformer 的 GNN 中的过度平滑与消息丰富度爆炸问题。
• 通过显式建模节点与边特征之间的交互，提升分子表示学习能力。
• 通过尊重分子拓扑结构的扩散机制整合结构归纳偏置，提升泛化能力与表达能力。

提出的方法

• CoMPT 采用基于 Transformer 的架构，通过节点与边之间的交叉注意力机制，迭代更新节点和边的嵌入表示。
• 引入基于分子图拓扑邻接矩阵的消息扩散机制，以约束消息传播，减少爆炸性消息累积。
• 使用可学习的节点位置和边方向嵌入，以保留结构上下文信息，提升表示质量。
• 消息传递以通信方式执行，节点与边以双向方式交换信息，增强特征表达能力。
• 通过可学习的读出机制和节点级回归头，框架支持图级别与节点级别的预测任务。
• 根据任务类型，使用交叉熵损失或平均绝对误差（MAE）损失进行端到端训练。

实验结果

研究问题

• RQ1 一种显式建模节点-边交互的基于 Transformer 的 GNN 是否能在分子表示学习中超越标准 GNN 和原始 Transformer？
• RQ2 引入拓扑消息扩散机制对分子图任务中模型性能与稳定性有何影响？
• RQ3 CoMPT 在多层架构中在多大程度上缓解了过度平滑问题并保持了可区分的原子表示？
• RQ4 节点位置、边方向与消息扩散机制对最终性能的贡献分别是什么？

主要发现

• 在七个图级别分子性质预测数据集上，CoMPT 相较于 SOTA 基线模型平均性能提升约 4%。
• 在节点级别任务中，CoMPT 在 1H-NMR 预测中达到 0.214 ppm 的平均绝对误差（MAE），在 13C-NMR 中为 1.321 ppm，优于所有基线模型。
• 消融实验表明，移除消息扩散机制导致性能下降最大，凸显其在控制消息爆炸中的关键作用。
• t-SNE 可视化结果证实，CoMPT 有效缓解了过度平滑问题，并实现了骨架原子与不同侧链官能团的清晰分离。
• 引入节点位置与边方向嵌入进一步提升了性能，表明其在结构感知中的重要性。
• 模型在多次随机种子实验中表现出强鲁棒性，重复实验的标准差较低（如 1H-NMR MAE 的标准差为 0.003）。