[論文レビュー] Don't Unroll Adjoint: Differentiating SSA-Form Programs

Michael Innes|arXiv (Cornell University)|Oct 18, 2018

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

ひとこと要約

本論文は SSA-form プログラムを用いた逆モード自動微分（AD）アプローチを提示し、Julia の Zygote として実装され、複雑な言語機能の効率的な微分とコンパイラとの相互作用を可能にする。

ABSTRACT

This paper presents reverse-mode algorithmic differentiation (AD) based on source code transformation, in particular of the Static Single Assignment (SSA) form used by modern compilers. The approach can support control flow, nesting, mutation, recursion, data structures, higher-order functions, and other language constructs, and the output is given to an existing compiler to produce highly efficient differentiated code. Our implementation is a new AD tool for the Julia language, called Zygote, which presents high-level dynamic semantics while transparently compiling adjoint code under the hood. We discuss the benefits of this approach to both the usability and performance of AD tools.

研究の動機と目的

• トレースベースの方法を超える表現力豊かで高性能な AD の必要性を動機づける。
• 既存のコンパイラと統合して最適化された微分コードを生成する SSA-form ベースの AD を提案する。
• Zygote の実装とそれが Julia のコンパイラとランタイムとどのように相互作用するかを示す。
• このアプローチが制御フロー・変異・高階微分といった言語機能をどのようにサポートするかを示す。

提案手法

• SSA-form IR を通して勾配を伝播させるために J 演算子と pullback を用いて微分を表現する。
• 逆方向の基本ブロックと phi ノードを用いて、随伴での制御フローとデータ依存性を処理する。
• 逆微分時の SSA イテレーション定義を管理するためにスタック上で alpha 値を格納・解決する。
• cons cell・ボックス・クロージャなどの複雑なデータ構造と変異を、それに対応する随伴 pullback で扱う。
• Julia の型推論とコンパイラ最適化を活用して pullbacks をインライン化し、効率的なコードを生成する。
• 高レベル微分可能プログラミングと従来のコンパイラ最適化の橋渡しとしてのアプローチを位置づける。

実験結果

リサーチクエスチョン

• RQ1 SSA-form IR 上の逆モード AD は制御フロー・変異・高階関数を含む完全な言語機能を収容できるか？
• RQ2 SSA-form プログラムを微分することが、ホストコンパイラとどのように相互作用して効率的な微分コードを生み出すか？
• RQ3 トレースベースの AD システムと比較した場合の性能と使いやすさのトレードオフは何か？
• RQ4 分岐とループを横断して正しいデータフローを保持するように随伴コードを生成するにはどうすればよいか？
• RQ5 このアプローチは実世界のモデルやコンパイラバックエンド（例：LLVM、TPU）に対してスケーラブルか？

主な発見

• Zygote は Julia 向けの SSA-form AD ツールとして機能し、高効率な微分コードを生み出す。
• 随伴の構築は phi ノードと alpha トラッキングを用いて、逆向き CFG を通じて正しいデータフローを保持する。
• Pullbacks は組み合わせ可能で、高階微分や変異・データ構造を含む複雑な言語機能をサポートできる。
• いくつかのケースで手書き微分と競合する性能を示し、トレース法に対して有利なベンチマークを示す。
• このアプローチは表現力を損なうことなく、既存のコンパイラ最適化や今後のアクセラレータを活用可能にする。
• SSAベースの微分は現代のコンパイラスタックと統合でき、言語横断の微分可能なプログラミングを可能にする。