[論文レビュー] Differentiating a Tensor Language

Gilbert Bernstein, Michael W. Mara+3 more|arXiv (Cornell University)|Aug 25, 2020

Parallel Computing and Optimization Techniques参考文献 29被引用数 5

ひとこと要約

この論文は、アイバーソンの指示関数を用いてスパarsityを明示的にモデル化することで、テンソルプログラムの逆方向誤差伝搬を完全に関数型で、かつ証明可能な効率性をもって計算する手法を提示する。新しい「テンソルSSA」正規形と内積に基づく自動微分の導出を活用することで、ミューテーションを避けながらも、コストの低い勾配の原則を達成し、導関数コードが元の関数と漸近的に同等の効率性を持つことを保証する。

ABSTRACT

How does one compile derivatives of tensor programs, such that the resulting code is purely functional (hence easier to optimize and parallelize) and provably efficient relative to the original program? We show that naively differentiating tensor code---as done in popular systems like Tensorflow and PyTorch---can cause asymptotic slowdowns in pathological cases, violating the Cheap Gradients Principle. However, all existing automatic differentiation methods that guarantee this principle (for variable size data) do so by relying on += mutation through aliases/pointers---which complicates downstream optimization. We provide the first purely functional, provably efficient, adjoint/reverse-mode derivatives of array/tensor code by explicitly accounting for sparsity. We do this by focusing on the indicator function from Iverson's APL. We also introduce a new "Tensor SSA" normal form and a new derivation of reverse-mode automatic differentiation based on the universal property of inner-products.

研究の動機と目的

• テンソルプログラムにおけるナーブな自動微分の非効率性、特に勾配計算における漸近的遅延を解消すること。
• 代入演算子（+=）のような命令型の副作用を避けることで、データ並列性を保持し、後続の最適化を可能にすること。
• 密度的およびスパースなテンソルに対して、完全に関数型の文脈でコストの低い勾配の原則を形式的に保証すること。
• 効率的かつ正しく導関数をコンパイル可能にするために、新しい正規形「テンソルSSA」を導入すること。
• 内積の普遍的性質に基づく、逆方向ADの数学的に厳密な基盤を提供すること。

提案手法

• ブール述語を表現し、テンソル演算におけるスパarsityを明示的にモデル化するために、アイバーソンの指示関数を用いる。
• 静的解析と最適化を可能にするために、新しい「テンソルSSA」正規形を導入する。
• 内積の普遍的性質を用いて導関数を導出することで、正しさと効率性を保証する。
• 随伴計算におけるスパarsityを考慮し、不要なメモリ書き込みを回避する。
• 漸近的効率性を保つ代数的簡略化を適用し、極端なケースでも有効である。
• スパarsityを一等級の概念として扱うことで、+= ミューテーションに依存せず、完全に関数型のコンパイルを可能にする。

実験結果

リサーチクエスチョン

• RQ1 命令型の副作用に依存せずに、完全に関数型のテンソル言語でコストの低い勾配の原則を達成できるか？
• RQ2 テンソル演算におけるスパarsityをどのように形式的にモデル化すれば、効率的な導関数計算を保証できるか？
• RQ3 どのような正規形が、データ並列的テンソルプログラムにおける効率的かつ正しく導関数を計算する逆方向自動微分を可能にするか？
• RQ4 内積の普遍的性質を用いて、証明可能な効率性を持つADシステムを導出できるか？
• RQ5 現在のフレームワーク（例：PyTorch、TensorFlow）における漸近的非効率性を体系的に回避できるか？

主な発見

• 提案手法は、完全に関数型の文脈でコストの低い勾配の原則を達成し、導関数計算が元のプログラムの定数倍以内の遅延で実行されることを保証する。
• アイバーソンの指示関数によるスパarsityのモデル化により、diag や trace 演算のような随伴計算において不要なメモリ書き込みが回避される。
• 内積および合計の調整を考慮しても、前向きパスと逆向きパスのコストが等価であることを証明する。
• テンソルSSA正規形は静的解析と最適化を可能にし、データ並列性を保持するとともに、効率的なコンパイルを可能にする。
• 極端なケースにおいても、ナーブな微分よりも性能が優れる。例えば、trace(A) + ... + trace(A) の勾配は、最適化された場合、O(kN²) ではなく O(kN) に低下する。
• += ミューテーションに依存しないため、関数型の純粋性が保たれ、より強力なコンパイラ最適化が可能になる。