[论文解读] Transformers over Directed Acyclic Graphs

Yuankai Luo, Thost, Veronika+1 more|arXiv (Cornell University)|Oct 24, 2022

Advanced Graph Neural Networks被引用 4

一句话总结

本文提出DAGFormer，一种通过新型注意力机制适配有向无环图（DAGs）的Transformer架构，该机制将感受野限制在可到达的前驱和后继节点上，从而在保持结构保真度的同时降低计算复杂度。此外，该方法还引入基于深度的位置编码以捕捉DAG的拓扑结构。在源代码和引用网络等多种DAG任务中，该方法在准确率和效率方面均优于标准图Transformer和DAG专用GNN。

ABSTRACT

Transformer models have recently gained popularity in graph representation learning as they have the potential to learn complex relationships beyond the ones captured by regular graph neural networks. The main research question is how to inject the structural bias of graphs into the transformer architecture, and several proposals have been made for undirected molecular graphs and, recently, also for larger network graphs. In this paper, we study transformers over directed acyclic graphs (DAGs) and propose architecture adaptations tailored to DAGs: (1) An attention mechanism that is considerably more efficient than the regular quadratic complexity of transformers and at the same time faithfully captures the DAG structure, and (2) a positional encoding of the DAG's partial order, complementing the former. We rigorously evaluate our approach over various types of tasks, ranging from classifying source code graphs to nodes in citation networks, and show that it is effective in two important aspects: in making graph transformers generally outperform graph neural networks tailored to DAGs and in improving SOTA graph transformer performance in terms of both quality and efficiency.

研究动机与目标

• 为解决标准Transformer在应用于DAG时缺乏结构归纳偏置的问题，DAG编码了部分顺序与可达性。
• 设计一种注意力机制，在保持表达能力的同时降低Transformer典型的二次方复杂度。
• 引入反映DAG部分序中节点深度的位置编码，以增强结构感知能力。
• 在多种DAG类型上评估该框架，证明其在准确率和推理速度方面优于GNN和现有图Transformer。
• 提供一种通用、高效且可扩展的框架，兼容现有Transformer与GNN架构。

提出的方法

• 提出一种DAG感知的注意力机制，将每个token的注意力限制在仅通过有向路径可达的节点上，有效将感受野限制在前驱和后继节点内。
• 使用可学习的掩码矩阵强制执行此限制，确保仅在拓扑有序且可达的节点上计算注意力。
• 引入基于深度的位置编码，显式编码每个节点在DAG部分序中的位置，以补充注意力机制。
• 将注意力计算公式化为可微分的稀疏操作，保持并行性，并相比标准自注意力显著降低FLOPs。
• 通过将框架设计为即插即用模块，实现与独立Transformer和消息传递GNN的集成。
• 以模块化方式实现该方法，支持部署在现有模型（如普通Transformer或SOTA图Transformer）之上。

实验结果

研究问题

• RQ1 能否有效将基于Transformer的架构适配以利用DAG的部分顺序与可达性结构？
• RQ2 仅将注意力限制在可达节点上，是否能同时提升DAG任务的模型效率与性能？
• RQ3 基于深度的位置编码在多大程度上增强了Transformer对DAG结构的表征能力？
• RQ4 所提出的框架能否在准确率和推理速度方面均优于标准图Transformer与DAG专用GNN？
• RQ5 在不同DAG数据集上，感受野大小（k）对性能与训练时间的影响如何？

主要发现

• 所提出的DAG注意力机制显著降低了计算复杂度，在ogbg-code2上实现每轮10分钟的训练时间，相比基线GNN（如DAGNN）的100分钟，速度提升达4倍。
• 在ogbg-code2数据集上，该方法在DAG+SAT设置下取得0.2018的测试F1分数，优于基线SAT模型（0.2000）和DAGNN（0.1975）。
• 将感受野大小（k）从1增加到∞，性能持续提升：DAG+TF的F1从0.1724提升至0.1879，DAG+SAT的F1从0.1533提升至0.2018。
• 在ogbg-code2上，当k=∞时，平均可达节点数（n_k）为9.78，远低于总节点数，表明即使在完全可达的情况下，该方法仍保持高效。
• 该框架在性能与效率方面均提升了SOTA图Transformer，证明通过注意力与位置编码引入结构偏置对DAG表征学习至关重要。
• 该方法泛化能力强，显著提升了多种DAG类型的表现，包括源代码图、引用网络与逻辑公式。