[論文レビュー] REGAL: Transfer Learning For Fast Optimization of Computation Graphs.

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

Parallel Computing and Optimization Techniques参考文献 20被引用数 12

ひとこと要約

REGAL は、実行時間とピークメモリ使用量を削減するために、ニューラルネットワークの計算グラフを最適化するための転移学習アプローチを提案する。最適化器は事前にオフラインで訓練され、未観測のグラフに一般化可能であり、従来の手法が数時間もかかるのに対し、数秒で高品質な最適化を達成する。これは著しく高速であり、古典的手法および学習ベースのベースラインを上回る性能を発揮する。

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つの訓練済み最適化器を多様な実世界のTensorFlowグラフに一般化すること。

提案手法

• 多様な計算グラフ上で事前にオフラインで訓練された深層強化学習エージェントが、最適化方針を学習する。
• エージェントは演算子の統合、レイアウト変換、メモリ割り当てなどの意思決定を学習する。
• 同じ方策ネットワークが、追加のトレーニングなしに未確認のグラフに一般化可能となるように、転移学習が適用される。
• 最適化意思決定を誘導するために、実行時間とメモリ使用量を組み合わせた報酬関数が使用される。
• サンプル効率を向上させるために、経験再生を用いたオフポリシー強化学習が方策ネットワークの訓練に用いられる。
• 意思決定を推論時に行うために、最適化コンパイラーパイプラインに統合される。

実験結果

リサーチクエスチョン

• RQ1 一度訓練された強化学習方策を、未確認の計算グラフに一般化して最適化できるか？
• RQ2 実行時間とメモリ使用量の観点から、提案手法は古典的手法と比べてどのように差がつくか？
• RQ3 転移学習により、新しいモデルを最適化する際の再トレーニングの必要性はどの程度減少するか？
• RQ4 最適化時間に顕著な高速化を達成しつつ、最適化品質を維持または向上させられるか？

主な発見

• REGAL は、実行時間とピークメモリ使用量の両面で、古典的手法および学習ベースのベースラインと同等またはそれ以上の最適化品質を達成する。
• 実世界のTensorFlowグラフにおいて、最適化時間を数時間から数秒に短縮する。
• 微調整なしに未確認のグラフに効果的に一般化され、優れた転送性が示された。
• 多様なモデルにおいて、実行時間とメモリフットプリントの両方を最小化する点で、強力なベースラインを上回る。
• 報酬ベースのRLフレームワークにより、速度とメモリ効率のトレードオフが効果的にバランスされた。