[論文レビュー] Dynamic Automatic Differentiation of GPU Broadcast Kernels

Jarrett Revels, Tim Besard+3 more|arXiv (Cornell University)|Oct 18, 2018

Parallel Computing and Optimization Techniques参考文献 22被引用数 4

ひとこと要約

本論文は、GPU最適化されたブロードキャスト演算を効率的に微分できるように、前向きモードと後向きモードの自動微分（AD）を組み合わせたミックスドモードAD手法を提案する。本手法は、ブロードキャストのヤコビ行列のスパarsityを活用し、ブロードキャストされた部分グラフにおける後向きモードのバックプロpagationを回避することで、データ依存の制御フローを伴う微分を効率的に行える。GPUワークロードにおいて、特にHM-LSTMセルの計算において、純粋な後向きモードアプローチを上回る性能を発揮する。

ABSTRACT

We show how forward-mode automatic differentiation (AD) can be employed within larger reverse-mode computations to dynamically differentiate broadcast operations in a GPU-friendly manner. Our technique fully exploits the broadcast Jacobian's inherent sparsity structure, and unlike a pure reverse-mode approach, this "mixed-mode" approach does not require a backwards pass over the broadcasted operation's subgraph, obviating the need for several reverse-mode-specific programmability restrictions on user-authored broadcast operations. Most notably, this approach allows broadcast fusion in primal code despite the presence of data-dependent control flow. We discuss an experiment in which a Julia implementation of our technique outperformed pure reverse-mode TensorFlow and Julia implementations for differentiating through broadcast operations within an HM-LSTM cell update calculation.

研究の動機と目的

• GPU上でデータ依存の制御フローを伴うブロードキャスト演算を処理する際、純粋な後向きモード自動微分の限界を解消すること。
• 前向きモードADを後向きモード計算内に統合することで、メモリおよび計算オーバーヘッドを低減する手法を開発すること。
• 純粋な後向きモードフレームワークでは通常制限されるが、データ依存の制御フローが存在する状況でも、プライマルコードにおけるブロードキャスト統合を可能にすること。
• ミックスドモードADが、純粋な後向きモード実装と比較して、GPUワークロードにおいて優れた性能を発揮することを示すこと。
• 局所的な構造に基づいて最適なADモードを動的に選択する次世代の微分コンパイラの設計を促すこと。

提案手法

• 本手法は、ブロードキャスト演算のヤコビアンに内在するスパarsity構造を活用することで、前向きモードADを後向きモードADとインタリーブ可能にする。
• 前向きモードでは多次元の双数（dual numbers）を用いて方向微分を計算し、ブロードキャスト関数の入力アリティに応じて摂動ベクトルをスケーリングする。
• 中間値の保存を避けることで、後向きモードのバックプロパゲーションに特有の制約を排除し、直接的に前向きモードの微分を計算する。
• JuliaのJITコンパイルパイプラインに動的にコード変換パスを挿入することで、コアコンパイラを変更せずに文脈に応じたAD動作を実現する。
• 制御フローが入力に依存する場合でも、ブロードキャスト演算の統合をサポートする。これは、完全なバックワードパスが必須であるため、純粋な後向きモードでは実現不可能である。
• ReverseDiff、Flux、Zygoteなどのパッケージを通じてJuliaのエコシステムに統合され、TPUのサポートも拡張済みである。

実験結果

リサーチクエスチョン

• RQ1 前向きモードADを後向きモードADと効果的にインタリーブすることで、純粋な後向きモードと比較して、GPU最適化ブロードキャスト演算の微分をより効率的に行えるか？
• RQ2 ブロードキャストヤコビアンのスパarsityを活用することで、後向きモードのバックプロパゲーションのオーバーヘッドを回避しつつ、性能向上が達成できるか？
• RQ3 本ミックスドモード手法は、通常は純粋な後向きモードADで禁止されるデータ依存の制御フローが存在する状況でも、ブロードキャスト統合をサポートできるか？
• RQ4 GPUアーキテクチャ上で、前向きモード微分における入力アリティの増加に伴い、本手法の性能はどのようにスケーリングするか？
• RQ5 実世界のディープラーニングワークロード、特にHM-LSTMセルの更新において、本手法は純粋な後向きモード実装をどの程度上回るか？

主な発見

• 本ミックスドモードADアプローチは、GPU上でHM-LSTMセルの更新計算におけるブロードキャスト演算の微分において、純粋な後向きモードのTensorFlowおよびJulia実装を上回った。
• 制御フローが入力に依存する状況でも、プライマルコードにおけるブロードキャスト統合が可能である。これは、純粋な後向きモードADではバックプロパゲーションの制約により実現不可能である。
• GPU上で高アリティのブロードキャスト演算を前向きモードで微分すると、入力アリティが5を超えると、レジスタ圧力の増加により、スレッドのスケジューリング効率とハードウェア利用効率が低下する。
• 18個の入力引数で、計算および帯域幅の利用効率が60％未満に低下し、スループットが低くなった。これは、スレッドのスループットが低いため、カーネルがレイテンシ制限に陥ったことを示している。
• 本技術は、FluxやZygoteなどの生産用途のJuliaベースMLフレームワークですでに導入されており、実用的かつ性能向上の有効性が裏付けられている。
• 本手法はTPUのサポートにも拡張されており、GPUアーキテクチャを越えた幅広い応用可能性を示している。