[論文レビュー] Demystifying Differentiable Programming: Shift/Reset the Penultimate Backpropagator

Fei Wang, Zheng, Daniel+4 more|arXiv (Cornell University)|Mar 27, 2018

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

ひとこと要約

本稿では、区切り付き継続（shift/reset）を用いた、逆モード自動微分（AD）の新規で簡素化された定式化を提案する。これにより、補助データ構造を必要とせず、演算子オーバーロードによる効率的で表現力豊かな微分可能プログラミングが可能になる。この手法をマルチステージプログラミングと組み合わせることで、TensorFlowの計算グラフの効率性とPyTorchの命令的スタイルの柔軟性を統合した、高性能なディープラーニングフレームワークLanternを構築した。

ABSTRACT

Deep learning has seen tremendous success over the past decade in computer vision, machine translation, and gameplay. This success rests in crucial ways on gradient-descent optimization and the ability to learn parameters of a neural network by backpropagating observed errors. However, neural network architectures are growing increasingly sophisticated and diverse, which motivates an emerging quest for even more general forms of differentiable programming, where arbitrary parameterized computations can be trained by gradient descent. In this paper, we take a fresh look at automatic differentiation (AD) techniques, and especially aim to demystify the reverse-mode form of AD that generalizes backpropagation in neural networks. We uncover a tight connection between reverse-mode AD and delimited continuations, which permits implementing reverse-mode AD purely via operator overloading and without any auxiliary data structures. We further show how this formulation of AD can be fruitfully combined with multi-stage programming (staging), leading to a highly efficient implementation that combines the performance benefits of deep learning frameworks based on explicit reified computation graphs (e.g., TensorFlow) with the expressiveness of pure library approaches (e.g., PyTorch).

研究の動機と目的

• プログラミング言語研究者を対象に、逆モード自動微分（AD）を、区切り付き継続との関連を明らかにすることで、解明すること。
• 明示的なテープやトレースデータ構造を必要とせず、shift/reset制御演算子を用いることで、逆モードADにおけるそれらの必要性を排除すること。
• マルチステージプログラミングを用いて、定義後に実行するフレームワーク（例：TensorFlow）のパフォーマンスと、定義しながら実行するフレームワーク（例：PyTorch）の表現力の両方を統合すること。
• 純粋にローカルな、演算子オーバーロードに基づく変換を用いて、再帰や高階関数を含む任意のプログラムにおける勾配計算を可能にすること。
• 実行時において計算グラフを再帰的に再構築するが、制限のない制御フローをサポートする、高性能なディープラーニングDSL、Lanternを構築すること。

提案手法

• CPS変換されたプログラムの記号的微分として逆モードADを形式化し、区切り付き継続との深い関連を明らかにする。
• shift/reset制御演算子を用いて逆モードADを実装し、補助テープを管理することなく、純粋にローカルで演算子オーバーロードに基づく微分が可能になるようにする。
• 継続に基づくAD定式化をマルチステージプログラミング（ステージング）と統合し、実行時に効率的で再帰的な計算グラフを生成する。
• ScalaのLightweight Modular Staging（LMS）フレームワークを用いて、ステージングと継続に基づくADをシームレスに組み合わせ、DSLの構築を可能にする。
• 得られたDSLをC++/CUDAにコンパイルし、高いパフォーマンスでの実行を実現し、推論および学習速度において競争力のある性能を達成する。
• ネストされたshift/resetを用いて多段階CPS変換をエンコードし、複雑な制御フローやネストされた抽象化の正しい取り扱いを可能にする。

実験結果

リサーチクエスチョン

• RQ1 区切り付き継続を用いて、明示的なテープ管理を伴わずに、演算子オーバーロードによる純粋な逆モード自動微分が実現可能か？
• RQ2 定義しながら実行するフレームワーク（例：PyTorch）の表現力と、定義後に実行するフレームワーク（例：TensorFlow）のパフォーマンスを、1つのシステムで統合できるか？
• RQ3 継続は、再帰的および高階関数を含む任意の第一級計算における勾配計算を可能にする役割を果たすか？
• RQ4 制限のない制御フローをサポートしつつ、効率的なグラフレベルの最適化やコンパイルを可能にするディープラーニングフレームワークを構築できるか？
• RQ5 圏論や複雑なデータ構造に依存しない、より単純で基礎的な逆モードADの定式化は存在するか？

主な発見

• 逆モードADは、shift/resetを用いた演算子オーバーロードにより、明示的なテープやトレースデータ構造を必要とせず実装可能である。
• 逆モードADと区切り付き継続との関連は、微分可能プログラミングのための洗練された、ローカルで合成可能な基盤を提供する。
• 継続に基づくADとマルチステージプログラミングを組み合わせることで、動的制御フローと効率的なグラフレベルのコンパilationを両立するフレームワークが実現できる。
• 得られたフレームワークLanternは、SqueezeNet、ResNet、DeepSpeech2、TreeLSTMといった最先端のモデルで競争力のあるパフォーマンスを達成し、既存のフレームワークを同等または上回る速度と効率性を示した。
• このアプローチは、一般の再帰や高階関数をネイティブにサポートでき、従来の明示的データ構造に依存するADシステムの主要な制限を克服した。
• 実装は、ディープラーニングフレームワークが、2つの直交的な抽象化であるステージングと区切り付き継続の合成として理解できると示した。