GO：グラフ最適化用強化学習(1/3)

１．GO：グラフ最適化用強化学習(1/3)まとめ

・大規模モデルは多様なアクセラレータが混在する環境でトレーニングされるようになってきた
・MLコンパイラは多くの複雑な最適化問題を解決する必要があるが手動設計には限界がある
・この制限を克服するため計算グラフ最適化用の深層強化学習であるGOを発表した

２．グラフ最適化用強化学習とは？

以下、ai.googleblog.comより「End-to-End, Transferable Deep RL for Graph Optimization」の意訳です。元記事の投稿は2020年12月17日、Yanqi ZhouさんとSudip Royさんによる投稿です。

グラフっぽい橋をGOしているアイキャッチ画像のクレジットはPhoto by Max Bender on Unsplash

ますます多くのアプリケーションが、多様なアクセラレータを使ってトレーニングされた大規模で複雑なニューラルネットワークによって実行されています。

このプロセスは、人間にとって理解しやすい計算グラフ等の表現を、機械にとって理解しやすいデバイス固有の実行可能ファイルに変換する機械学習(ML:Machine Learning)用コンパイラを使うと容易になります。

その際、MLコンパイラは、グラフの書き換え、デバイス配置の問題、複数操作の融合、行列のレイアウト修正や充填、スケジューリングなど、多くの最適化問題を解決する必要があります。

例えば、デバイス配置の問題(訳注： device placement problem。CPUやGPUが多数混在している環境でどの処理を何で行うのが最も効率的かを知る事)では、コンパイラは、トレーニングステップ時間などの目的関数を最小化できるように、計算グラフの操作とターゲットとなる物理デバイスへの割り当てを決定する必要があります。配置のパフォーマンスは、デバイス間のネットワーク帯域幅、デバイスの最大メモリ、同時配置制約などの複雑な要因の組み合わせによって決定されます。

そのため、経験則または探索ベースのアルゴリズムで解決するのは困難になります。これらのアルゴリズムは通常、高速ですが最適ではないソリューションに落ち着きます。更に、経験則は、特に新しいモデルアーキテクチャが出現した際に、新たに開発および維持するのが困難です。

学習ベースのアプローチを使用する最近の試みは、有望な結果を示していますが、実際に製品展開する事が不可能になるいくつかの制限があります。

第1に、これらのアプローチは、特に新しいモデルアーキテクチャから生じる、初めてみるグラフに簡単に一般化できません。第2に、サンプル効率が低く、トレーニング中のリソース消費が高くなります。最後に、それらは単一の最適化タスクを解決することしかできず、その結果、コンパイルスタック内の密に結合した最適化問題全体の依存関係を捕捉できません。

直近のNeurIPS 2020で口頭論文として発表された「Transferable Graph Optimizers for ML Compilers」では、上記の全ての制限を克服する計算グラフ最適化(GO:Graph Optimization)のためのエンドツーエンドで転移可能な深層強化学習を提案します。

TensorFlowのデフォルトの最適化と比較して、3つのグラフ最適化タスクで33％～60％のスピードアップを示します。

Inception-v3、Transformer-XL、WaveNetなど、最大80,000ノードで構成される代表的なグラフの多様なセットで、GOは、人間のエキスパートによる最適化よりも平均21%向上し、従来の最先端技術よりも18%向上し、収束が15倍速くなります。

MLコンパイラのグラフ最適化問題
MLコンパイラには頻繁に発生する3組の最適化タスクがあります。これらを学習したポリシーを使用して解決できる決定問題として定式化します。各タスクの決定問題は、「計算グラフの各ノードの決定を行う事」と再構成できます。

最初の最適化タスクはデバイス配置の問題です。ここでの目標は、エンドツーエンドの実行時間が最小になるように、グラフのノードを実行する物理デバイスに割り当てる最適な方法を決定する事です。

2番目の最適化タスクは演算操作のスケジューリングです。計算グラフの演算操作は、入力テンソルがデバイスメモリに存在するときに実行する事ができます。良く使われるスケジューリング戦略は、各デバイスで操作準備ができたキューを維持し、先入れ先出しの順序で演算操作をスケジュールすることです。ただし、このスケジューリング戦略で、他のデバイスに配置された下流操作は考慮されていないのでブロックが発生する可能性がある事と、多くの場合、十分に活用されていないデバイスでスケジュールが作成されます。

このようなデバイス間の依存関係を追跡できるスケジュールを見つけるために、私達のアプローチでは、優先度ベースのスケジューリングアルゴリズムを使用して、それぞれの優先度に基づいて準備完了キュー内の演算操作をスケジュールします。デバイスの配置と同様に、演算操作のスケジューリングは、「実行時間に基づいて報酬を最大化するためにグラフ内の各ノードに優先順位を割り当てるポリシーを学習する問題」として定式化できます。

3番目の最適化タスクは複数操作の融合(operation fusion)です。話を簡潔にするために、ここではこの問題の詳細な説明を省略します。代わりに、優先度ベースのスケジューリングと同様に、複数操作の融合でも優先度ベースのアルゴリズムを使用して、融合するノードを決定できることに注意してください。この場合のポリシーネットワークの目標は、グラフ内の各ノードに優先順位を割り当てることです。

最後に、3つの最適化問題のそれぞれで行われた決定が、他の問題の最適決定に影響を与える可能性がある事を認識することが重要です。例えば、2つのノードを2つの異なるデバイスに配置すると、操作の融合は事実上無効になり、スケジューリングに影響を与える可能性のある通信遅延が発生します。