[論文レビュー] Reinforcement Learning and Adaptive Sampling for Optimized DNN Compilation

Byung Hoon Ahn, Prannoy Pilligundla+1 more|arXiv (Cornell University)|May 30, 2019

Advanced Neural Network Applications参考文献 65被引用数 13

ひとこと要約

本論文では、クラスタリングを用いた適応的サンプリングを統合することで、高価なハードウェア測定を削減する強化学習（RL）ベースの最適化コンパイラReLeASEを提案する。DNNコンパイルをRL問題として定式化し、代表的な部分空間にサンプルを集中させることで、ReLeASEはAutoTVMと比較して最適化時間を4.45倍速くし、推論性能を最大5.6%向上させた。また、探索ステップ数と測定オーバーヘッドが顕著に削減された。

ABSTRACT

Achieving faster execution with shorter compilation time can enable further diversity and innovation in neural networks. However, the current paradigm of executing neural networks either relies on hand-optimized libraries, traditional compilation heuristics, or very recently, simulated annealing and genetic algorithms. Our work takes a unique approach by formulating compiler optimizations for neural networks as a reinforcement learning problem, whose solution takes fewer steps to converge. This solution, dubbed ReLeASE, comes with a sampling algorithm that leverages clustering to focus the costly samples (real hardware measurements) on representative points, subsuming an entire subspace. Our adaptive sampling not only reduces the number of samples, but also improves the quality of samples for better exploration in shorter time. As such, experimentation with real hardware shows that reinforcement learning with adaptive sampling provides 4.45x speed up in optimization time over AutoTVM, while also improving inference time of the modern deep networks by 5.6%. Further experiments also confirm that our adaptive sampling can even improve AutoTVM's simulated annealing by 4.00x.

研究の動機と目的

• 現在、ヒューリスティック法や確率的法（例：模擬アニーリングや遺伝的アルゴリズム）に依存しているDNNコンパイル最適化の時間的・コスト的負担を軽減すること。
• 探索中に生じる高コストな実ハードウェア測定を最小限に抑えるために、必要なサンプル数を削減すること。
• 従来の探索戦略に代えて強化学習を活用することで、DNNコンパイラー最適化における探索効率と解の品質を向上させること。
• クラスタリングに基づく部分空間の集中を実現するサンプリング戦略を開発し、高い代表性と関連性を維持しながら測定回数を削減すること。

提案手法

• DNNコンパイルを強化学習（RL）問題として定式化し、エージェントが最適なコンパイラパラメータ（ノブ）を選択して性能を最大化するように学習する。
• 探索空間をクラスタリングすることで代表的な部分空間を同定し、これらの点にハードウェア測定を優先する適応的サンプリングアルゴリズムを導入する。
• クラスタリングを用いて部分空間全体を要約することで、高価な実ハードウェア測定の回数を削減しながらも、探索の質を保持する。
• 測定された推論時間に基づく報酬関数を用い、RLエージェントがより高速で最適化されたコードに到達するように誘導する。
• RLポリシー学習とコストモデルを組み合わせ、探索空間の情報量の高い領域に焦点を当てることで、重複する測定を回避する。
• ReLeASEとしてフレームワークを実装し、TVMおよびその他のDNNフレームワークと互換性を持つオープンソースの最適化コンパイラとして提供する。

実験結果

リサーチクエスチョン

• RQ1 強化学習は、模擬アニーリングや遺伝的アルゴリズムといった従来の確率的最適化手法よりも、DNNコンパイラー最適化において優れているか？
• RQ2 クラスタリングに基づく適応的サンプリングは、解の品質を損なわせることなく、必要なハードウェア測定回数を削減できるか？
• RQ3 強化学習と適応的サンプリングを組み合わせることで、実世界のDNNにおいて収束が速くなり、より最適化されたコードが得られるか？
• RQ4 適応的サンプリングは、もともとコストモデルを用いているAutoTVMのような既存のフレームワークに、どの程度向上効果をもたらすか？

主な発見

• ReLeASEはResNet-18において、AutoTVMと比較して最適化時間を4.45倍速くし、1.20時間（ReLeASE）対9.13時間（AutoTVM）を記録した。
• AlexNet、VGG-16、ResNet-18の全モデルで推論性能を最大5.6%向上させ、エンドツーエンドの推論時間を1.0277 msから0.9673 ms（AlexNet）に短縮した。
• 適応的サンプリングは、AutoTVMの模擬アニーリングを最適化時間で4.00倍速くし、汎用性を実証した。
• RLベースのアプローチは、模擬アニーリングよりも少ないステップで収束し、より良いまたは同等の性能を、少ない探索努力で達成した。
• 適応的サンプリングアルゴリズムは、クラスタリングされた代表的部分空間に焦点を当てることで、ハードウェア測定回数を削減し、サンプル効率を向上させた。
• ReLeASEはオープンソースとして公開されており、既存のDNNコンパイルパイプラインへの統合を可能にした。