[論文レビュー] Reinforced Genetic Algorithm Learning for Optimizing Computation Graphs

Aditya Paliwal, Felix Gimeno+5 more|arXiv (Cornell University)|May 7, 2019

Metaheuristic Optimization Algorithms Research参考文献 42被引用数 25

ひとこと要約

本稿では、実行時間とピークメモリ使用量を最小化するため、ニューラルネットワーク計算グラフの最適化を支援する強化学習を統合した遺伝的アルゴリズムであるREGALを提案する。BRKGA（バイアス付きランダムキー遺伝的アルゴリズム）の探索を、入力に依存する提案分布の学習によってガイドするGNN（グラフニューラルネットワーク）を訓練することで、再トレーニングを必要とせずに、未学習の実世界のTensorFlowグラフに対しても、数秒で高品質な解を得られる。従来の古典的手法および学習ベースのベースラインを上回る性能を発揮する。

ABSTRACT

We present a deep reinforcement learning approach to minimizing the execution cost of neural network computation graphs in an optimizing compiler. Unlike earlier learning-based works that require training the optimizer on the same graph to be optimized, we propose a learning approach that trains an optimizer offline and then generalizes to previously unseen graphs without further training. This allows our approach to produce high-quality execution decisions on real-world TensorFlow graphs in seconds instead of hours. We consider two optimization tasks for computation graphs: minimizing running time and peak memory usage. In comparison to an extensive set of baselines, our approach achieves significant improvements over classical and other learning-based methods on these two tasks.

研究の動機と目的

• 再トレーニングを必要とせず、多様な実世界のモデルに適用可能な高速で一般化可能なニューラルネットワーク計算グラフ最適化ツールの開発。
• 実行コスト（実行時間またはピークメモリ）を最小化する目的で、スタティックコンパイラにおける配置とスケジューリング意思決定を統合的に最適化すること。
• 1つのグラフあたり数時間のトレーニングを要するか、未学習のアーキテクチャへの一般化が不十分な、先行の学習ベース最適化手法の限界を克服すること。
• 大規模なグラフにおいても、数秒以内に高品質な解を得ることで、最適化コンパイラの本番環境での導入を可能にすること。

提案手法

• BRKGAの変異生成段階において、各ノードの非一様なベータ分布のパラメータを予測するため、REINFORCEを用いた文脈的バンドイット設定でGNNを訓練する。
• BRKGAは、一様なデフォルト値の代わりに、入力に依存する学習済みの提案分布を用いることで、より高品質な解に向けた探索をガイドする。
• GNNは、計算グラフの構造、ノード特徴量（例：演算タイプ、テンソルサイズ）、エッジ依存関係を処理し、配置およびスケジューリングのためのノード固有の行動確率を出力する。
• 報酬信号として、実行時間やメモリ使用量などの目的値をスパースな報酬として用い、明示的な教師信号は一切不要である。
• 候補解の実行コストを推定する粗い静的コストモデルを用いることで、BRKGAの評価段階で迅速なフィードバックが得られ、トレーニングの高速化を実現する。
• グラフのトポロジーの多様性、特に実世界のTensorFlowモデルを含む、多様なグラフ構造に対して、一般化可能な意思決定パターンを学習することで、汎用性を実現する。

実験結果

リサーチクエスチョン

• RQ1 学習済み方策は、再トレーニングを伴わず、未学習の計算グラフにおいても、より優れた解を遺伝的アルゴリズムが探索できるように導けるか？
• RQ2 GNNを用いた方策学習は、スタティックコンパイラにおける配置とスケジューリングの統合最適化において、解の品質を向上させるか？
• RQ3 本手法は、数時間ではなく数秒で高品質な結果を得られるか。これにより、本番環境のコンパイラにおけるリアルタイム利用が可能になるか？
• RQ4 本モデルは、多様なニューラルネットワークアーキテクチャおよびワークロードに一般化できるか？

主な発見

• REGALは、実世界のTensorFlowグラフにおいて、実行時間とピークメモリ使用量の両面で、古典的手法および学習ベースのベースラインを顕著に上回る。
• 先行の学習ベース手法（例：Mirhoseini et al. (2017, 2018)）が数時間のトレーニングを要するのに対し、REGALは数秒で高品質な解を生成できる。
• GNN方策は、未学習のグラフトポロジー、特にデータ拡張によって生成されたものに対しても、効果的に一般化され、異なるアーキテクチャ間で性能のばらつきが小さい。
• 方策は単純なヒューリスティクスをはるかに超える、非自明なデータ依存戦略を学習している。例として、高メモリノードに対してはスケジューリング優先度を低く設定し、重いノードに対してはデバイスバイアスを回避する戦略が含まれる。
• TF Runtimeテストセットにおいて、REGALは95%のケースでBRKGAベースラインを上回り、報酬が-1より大きいケースは、より優れた解を示している。