结构感知Transformer - 为图表示学习带来新视角

摘要： 结构感知Transformer（SAT）是一个基于PyTorch Geometric实现的创新性框架，它在标准Transformer的基础上融入了结构信息，通过提取节点为中心的子图表示来增强自我注意力机制。该项目提供了一种简单且灵活的方法，可将任何现有的图神经网络（GNN）用于子图表示，从而系统地提高与基础GNN相比的表现。

项目介绍

结构感知Transformer（SAT）是针对图表示学习的最新进展，由Dexiong Chen、Leslie O'Bray和Karsten Borgwardt在ICML 2022会议上提出。这个开源实现旨在将结构信息纳入到传统的Transformer架构中，以解决传统方法忽视图结构的问题。SAT的核心在于其自注意力层的"结构化"，它使用结构提取器（如GNN）来获取每个节点的子图表示，并以此计算注意力权重。

项目技术分析

SAT与标准Transformer的对比

SAT不同于标准Transformer之处在于，它不直接使用原始节点特征进行自注意力计算，而是先利用一个结构提取器（例如$k$-subtree或$k$-subgraph）获取以每个节点为中心的子图表示。然后，这些更新的节点嵌入被用于构建查询矩阵$\mathbf{Q}$和键矩阵$\mathbf{K}$，从而使得注意力机制能够考虑到局部结构信息。

结构提取器

SAT提供了两种结构提取器示例：$k$-subtree和$k$-subgraph。前者通过GNN处理以每个节点为中心的$k$-跳树，后者则对子图进行完全展开并使用GNN对整个子图进行建模。

快速上手

以下是在ZINC数据集上的快速启动代码示例：

# 省略...
dset = datasets.ZINC('./datasets/ZINC', subset=True, split='train')
dset = GraphDataset(dset)
train_loader = DataLoader(dset, batch_size=16, shuffle=True)
model = GraphTransformer(...)
output = model(data)
# 省略...

应用场景

• 化学分子预测：SAT适用于药物发现领域中的分子性质预测任务，例如ZINC数据集上的化学反应预测。
• 社交网络分析：理解用户之间的互动模式和社区结构。
• 代码理解：在OGBG-CODE2等数据集上进行编程语言结构的分类任务。
• 生物信息学：预测蛋白质相互作用和基因网络结构。

项目特点

1. 灵活性：可以结合任意GNN作为结构提取器。
2. 结构融合：将结构信息直接注入自注意力机制，提高了模型对图结构的理解。
3. 易于使用：提供清晰的API和示例脚本，便于快速集成到现有工作流中。
4. 可视化：内置模型解释功能，直观展示节点间注意力权重。

对于那些寻求更深入理解和改进图神经网络性能的研究者，结构感知Transformer无疑是一个值得探索的前沿工具。立即安装并开始您的图表示学习旅程吧！