[論文レビュー] ProTuner: Tuning Programs with Monte Carlo Tree Search

Ameer Haj-Ali|arXiv (Cornell University)|May 27, 2020

Machine Learning and Algorithms参考文献 47被引用数 7

ひとこと要約

ProTunerは、深層学習および画像処理ワークロード向けに、Halideプログラムの自動チューニングにモンテカルロ木探索（MCTS）を活用する。スケジューリングをマルコフ決定過程（MDP）として定式化し、前方探索と長期報酬推定に上界確信度（UCB）を用いたMCTSを活用することで、ビームサーチを凌駆し、16の実ベンチマークで最大3.25倍の性能向上を達成するとともに、コストモデルのノイズに対してより頑健である。

ABSTRACT

We explore applying the Monte Carlo Tree Search (MCTS) algorithm in a notoriously difficult task: tuning programs for high-performance deep learning and image processing. We build our framework on top of Halide and show that MCTS can outperform the state-of-the-art beam-search algorithm. Unlike beam search, which is guided by greedy intermediate performance comparisons between partial and less meaningful schedules, MCTS compares complete schedules and looks ahead before making any intermediate scheduling decision. We further explore modifications to the standard MCTS algorithm as well as combining real execution time measurements with the cost model. Our results show that MCTS can outperform beam search on a suite of 16 real benchmarks.

研究の動機と目的

• 手動チューニングが高コストであり、ヒューリスティクスが最適スケジュールを発見できない場合が多い、深層学習および画像処理ワークロードにおける自動プログラムスケジューリングの課題に対処すること。
• ビームサーチの欠点（貪欲性とノイズの多い中間的性能予測に依存すること）を克服し、より頑健で前方探索に基づく探索戦略を採用すること。
• MCTSとUCBを用いて、プログラムスケジューリング問題をマルコフ決定過程（MDP）として適用する可能性と有効性を検討すること。
• 学習済みコストモデルと実実行時間測定値を組み合わせた影響が、スケジューリング性能に与える影響を評価すること。
• MCTSにおける主な性能ボトルネック（例えば、子ノードの過剰な列挙オーバーヘッドやコストモデルへの過剰適合）を特定し、それらを是正すること。

提案手法

• 状態が中間スケジュールを表し、アクションがスケジューリング最適化（例：タイリング、ベクトル化、統合）であるマルコフ決定過程（MDP）としてHalideスケジューリング問題を定式化する。
• 探索と活用のバランスを図るため、上界確信度（UCB）を用いたモンテカルロ木探索（MCTS）を適用し、前方探索と長期報酬推定を可能にする。
• 完全なスケジュールの評価にシミュレーションベースのロールアウトを用いることで、誤差の多い中間的性能予測に依存しなくなるようにする。
• 学習済みコストモデルと実実行時間測定値を統合することで、探索の正確性を向上させ、モデルの不正確さへの過剰適合を低減する。
• 探索ハイパーパrameter $C_p$ のチューニングと、シミュレーション中の重複した子ノードの列挙を減らすことで、MCTSの性能を最適化する。
• 各意思決定後に新しいルートで探索木を再構築することで、適応的かつ動的なスケジューリング探索を可能にする。

実験結果

リサーチクエスチョン

• RQ1 MCTSにUCBを適用した場合、深層学習および画像処理ワークロード向けの高パフォーマンスHalideスケジュールをビームサーチを上回って発見できるか？
• RQ2 MCTSにおける完全なスケジュールシミュレーションによる前方探索は、ビームサーチと比較して、ノイズが多いまたは不正確なコストモデルの影響を軽減できるか？
• RQ3 実実行時間測定値と学習済みコストモデルを組み合わせることで、自動チューニングプロセスの性能と頑健性にどのような影響を与えるか？
• RQ4 MCTSベースの自動チューニングにおける主な性能ボトルネックは何か、そしてそれらはどのように緩和できるか？
• RQ5 異なるプログラムや探索深さにおいて、探索パrameter $C_p$ を動的にチューニングすることで、MCTSの効率性と有効性が向上するか？

主な発見

• ProTunerにMCTSを適用した場合、評価した16の実ベンチマークすべてでビームサーチを上回り、最大3.25倍の性能向上を達成した。
• MCTSにおける完全なスケジュールシミュレーションの活用により、ノイズの多い中間的コスト推定の影響が軽減され、ビームサーチよりも探索がより頑健になった。
• 実実行時間測定値と学習済みコストモデルを組み合わせることで、さらに性能が向上し、ハイブリッド評価の価値が示された。
• 現在のMCTS実装ではコスト評価に7.5％のオーバーヘッドが生じており、その88％が無駄な子ノードの列挙に費やされているため、最適化により最大8倍の性能向上が見込める。
• MCTSの実行時間が長すぎるとコストモデルへの過剰適合が発生するため、早期停止または動的探索終了戦略の導入が求められる。
• 異なるプログラムや探索深さにおいて $C_p$ を動的にチューニングすることで、分散が小さい短時間実行のプログラムにおいて特に探索効率が向上する可能性がある。