[論文レビュー] Transformers over Directed Acyclic Graphs

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

Advanced Graph Neural Networks被引用数 4

ひとこと要約

本稿では、有向無閉路グラフ（DAG）に適応されたトランスフォーマーモデルDAGFormerを提案する。新規のアテンション機構により、受容 field を到達可能な先行ノードおよび後続ノードに制限することで、計算複雑度を低減しながら構造の忠実性を維持する。さらに、DAGのトポロジーを捉えるために深さに基づく位置エンコーディングを統合する。本手法は、ソースコードや引用ネットワークを含む多様なDAGタスクにおいて、標準的なグラフトランスフォーマーや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.

研究の動機と目的

• DAGに適用された標準的なトランスフォーマーに構造的インダクティブバイアスが欠如していること、特に部分順序と到達可能性を捉えることの重要性を解決する。
• トランスフォーマー特有の2次関数的複雑度を低減しつつ表現力を維持するアテンション機構の設計。
• DAGの部分順序におけるノードの深さを反映する位置エンコーディングを統合し、構造的認識を強化する。
• 多様なDAGタイプにわたるフレームワークの評価を通し、GNNおよび既存のグラフトランスフォーマーを上回ることを示す。
• 既存のトランスフォーマーやGNNアーキテクチャと互換性があり、汎用的で効率的かつ拡張可能なフレームワークを提供すること。

提案手法

• 各トークンのアテンションを、有向パスによって到達可能なノードに制限するDAGに配慮したアテンション機構を導入する。これにより、受容 field が先行ノードおよび後続ノードに限定される。
• 制限を強制するために学習可能なマスク行列を用い、トポロジカルに順序付けられかつ到達可能なノードのみにアテンションが計算されるようにする。
• DAGの部分順序におけるノードの位置を明示的にエンコードする深さベースの位置エンコーディングを導入し、アテンション機構を補完する。
• アテンション計算を微分可能でスパースな操作として定式化し、標準的な自己アテンションと比較してFLOPsを削減しつつ並列性を維持する。
• プラグインモジュールとしての設計により、スタンドアロンのトランスフォーマーやメッセージパッシングGNNへの統合を可能にする。
• モジュラーな実装により、ヴァニラトランスフォーマーやSOTAグラフトランスフォーマーなどの既存モデルの上に容易に統合可能である。

実験結果

リサーチクエスチョン

• RQ1 部分順序と到達可能性というDAGの構造的特徴を効果的に活用できるように、トランスフォーマーベースのアーキテクチャを適応可能とすることができるか？
• RQ2 到達可能なノードにのみアテンションを制限することで、DAGタスクにおけるモデルの効率性と性能が向上するか？
• RQ3 深さベースの位置エンコーディングは、トランスフォーマーにおけるDAG構造表現の向上にどの程度寄与するか？
• RQ4 提案されたフレームワークは、標準的なグラフトランスフォーマーおよびDAG特化型GNNを上回る精度と推論速度を達成できるか？
• RQ5 受容 field のサイズ（k）を変化させた場合、異なるDAGデータセットにおいて性能と学習時間にどのような影響が生じるか？

主な発見

• 提案されたDAGアテンション機構により、計算複雑度が顕著に低減され、DAGNNなどのベースラインGNNと比較して最大4倍速い学習が達成された。ogbg-code2では1エポックあたり10分（DAGNNは100分）を記録した。
• ogbg-code2データセットにおいて、DAG+SATを用いた場合、テストF1スコアは0.2018を達成し、ベースラインのSATモデル（0.2000）およびDAGNN（0.1975）を上回った。
• 受容 field サイズ（k）を1から∞に増加させたところ、性能が一貫して向上し、DAG+TFではF1が0.1724から0.1879に、DAG+SATでは0.1533から0.2018に上昇した。
• ogbg-code2におけるk=∞の平均到達ノード数（n_k）は9.78と、全ノード数に比べて著しく低く、全到達可能性を考慮しても効率性が保たれていることが示された。
• フレームワークはSOTAグラフトランスフォーマーの品質と効率性を両方向上させ、アテンションと位置エンコーディングによる構造的バイアスがDAG表現学習において極めて重要であることを示した。
• 本手法は一般化能に優れ、ソースコードグラフ、引用ネットワーク、論理式など多様なDAGタイプにおいて性能向上を示した。