[論文レビュー] Neural Topological Ordering for Computation Graphs

Mukul Gagrani, Corrado Rainone+8 more|arXiv (Cornell University)|Jul 13, 2022

Advanced Graph Neural Networks被引用数 4

ひとこと要約

本論文では、コンパイラワークロードにおけるピークメモリ使用量を最小化するための、有向無閉路グラフ（DAG）における最適なトポロジカル順序を学習する、エンドツーエンドの深層学習フレームワーク、Topoformerを提案する。非自己回帰的で、アテンションベースのグラフニューラルネットワークを用い、トポロジカルに注意を払ったメッセージパッシングを実装することで、ベースラインと比較して著しく高速な推論を実現しながら、最先端または同等の性能を達成した。

ABSTRACT

Recent works on machine learning for combinatorial optimization have shown that learning based approaches can outperform heuristic methods in terms of speed and performance. In this paper, we consider the problem of finding an optimal topological order on a directed acyclic graph with focus on the memory minimization problem which arises in compilers. We propose an end-to-end machine learning based approach for topological ordering using an encoder-decoder framework. Our encoder is a novel attention based graph neural network architecture called \emph{Topoformer} which uses different topological transforms of a DAG for message passing. The node embeddings produced by the encoder are converted into node priorities which are used by the decoder to generate a probability distribution over topological orders. We train our model on a dataset of synthetically generated graphs called layered graphs. We show that our model outperforms, or is on-par, with several topological ordering baselines while being significantly faster on synthetic graphs with up to 2k nodes. We also train and test our model on a set of real-world computation graphs, showing performance improvements.

研究の動機と目的

• コンパイラパイプラインにおけるメモリ最小化を目的とした、DAGにおける最適なトポロジカル順序を求める組合せ最適化問題に対処すること。
• ヒューリスティック手法に代わる、学習された方策を用いるスケーラブルでエンドツーエンドの深層学習フレームワークを構築すること。
• DAG上でトポロジカルに注意を払ったグローバルなメッセージパッシングを可能にする、新しいグラフニューラルネットワークアーキテクチャ、Topoformerを導入すること。
• 再現可能なベンチマークのための、レイヤー構造を持ち、ニューラルネットワークに類似した計算グラフからなる、公開可能な合成データセットを構築すること。
• 合成および実世界の計算グラフの両方で、高速な推論速度を維持しながら性能向上を示すこと。

提案手法

• 非自己回帰的（NAR）デコーダーを採用することで、トポロジカル順序の確率分布を効率的にパrameter化し、高速な推論を実現する。
• エンコーダーとして、ノードペア間のトポロジカル関係（例：先祖、子孫、関係なし）に応じて異なる可学習パラメータを用いてメッセージパッシングを実行する、新しいアテンションベースのグラフニューラルネットワーク、Topoformerを採用する。
• Topoformerは、前方および後方エッジ伝搬を含む7種類の異なるトポロジカル変換を経て、DAG全体の構造的依存関係を捉える。
• Topoformerから得られるノード埋め込みは、デコーダーのトポロジカルシーケンス生成を導くノードの優先順位に変換される。
• モデルは、実際のニューラルネットワークリソースに類似した構造を持ち、スケーラブルなデータ生成が可能な、レイヤー構造の計算グラフからなる合成データセットで学習される。
• 教師あり学習を用い、近似動的計画法（DP）の解をターゲットとして、方策ネットワークを監視する。

実験結果

リサーチクエスチョン

• RQ1 学習可能でエンドツーエンドの深層学習フレームワークは、DAGにおけるピークメモリ使用量の最小化において、従来のヒューリスティック手法を上回ることができるか？
• RQ2 DAGにおけるGNNアーキテクチャにおいて、トポロジカルに注意を払ったグローバルなメッセージパッシングは、局所的またはエッジのみのメッセージパッシングと比較して、どの程度有効であるか？
• RQ3 非自己回帰的デコーディング方式は、自己回帰的ベースラインと比較して、著しく短い推論時間で競争力のある性能を達成できるか？
• RQ4 レイヤー構造を持つ、公開可能な合成DAGデータセットは、機械学習における組合せ最適化の分野で、再現可能なベンチマークをどの程度可能にするか？
• RQ5 提案手法は、合成グラフにとどまらず、実世界のコンパイラワークロードにも一般化可能か？

主な発見

• 500ノードの合成グラフでは、ビームサーチを用いた完全なTopoformerモデルが、近似DPの解から平均3.08 ± 0.51%のギャップを達成し、GAT（前方＋後方）および他のベースラインを上回った。
• 1000ノードのグラフでは、ビームサーチを用いたTopoformerが、0.05 ± 0.36%のギャップを達成し、GAT（前方＋後方）および他のモデルを著しく上回った。
• 完全なTopoformerモデルは、前方＋後方バージョンと比較して、近似DPからのギャップを1.5–2.5パーセンテージポイント低減した。これはグローバルなメッセージパッシングの利点を示している。
• Topoformerは、500ノードグラフにおけるビームサーチの平均推論時間1.62秒で最先端の性能を達成したが、ベースラインは最大4000倍も長い時間を要した。
• 実世界の計算グラフにおいても、強力なベースラインを常に上回る一貫した性能向上を示し、産業的コンパイラワークフローへの関連性を裏付けた。
• 提案された合成グラフ生成アルゴリズムは、スケーラブルで再現可能なベンチマークを可能にし、研究コミュニティ向けに公開可能である。