[論文レビュー] A Differentiable Programming System to Bridge Machine Learning and Scientific Computing

Mike Innes, Alan Edelman+5 more|arXiv (Cornell University)|Jul 17, 2019

Neural Networks and Applications被引用数 129

ひとこと要約

本論文は、Julia で構築された微分可能プログラミングシステム ∂P を提示する。これにより一般的な Julia プログラムを微分可能に通過させ、ML および科学計算パッケージと統合し、ディープラーニングから微分可能な物理学・金融まで多様な応用を実証する。

ABSTRACT

Scientific computing is increasingly incorporating the advancements in machine learning and the ability to work with large amounts of data. At the same time, machine learning models are becoming increasingly sophisticated and exhibit many features often seen in scientific computing, stressing the capabilities of machine learning frameworks. Just as the disciplines of scientific computing and machine learning have shared common underlying infrastructure in the form of numerical linear algebra, we now have the opportunity to further share new computational infrastructure, and thus ideas, in the form of Differentiable Programming. We describe Zygote, a Differentiable Programming system that is able to take gradients of general program structures. We implement this system in the Julia programming language. Our system supports almost all language constructs (control flow, recursion, mutation, etc.) and compiles high-performance code without requiring any user intervention or refactoring to stage computations. This enables an expressive programming model for deep learning, but more importantly, it enables us to incorporate a large ecosystem of libraries in our models in a straightforward way. We discuss our approach to automatic differentiation, including its support for advanced techniques such as mixed-mode, complex and checkpointed differentiation, and present several examples of differentiating programs.

研究の動機と目的

• ML と科学計算にまたがる統一的な微分可能プログラミングアプローチの提唱。
• 最小限のユーザー介入で一般的なプログラムを微分可能にできる Julia ベースのシステムを提供。
• 機械学習と既存の科学パッケージを組み合わせた実用的な応用を示す（例: レイトレーシング、量子シミュレーション、ファイナンス）。
• カスタムデータ型と複雑な制御フローによる性能特性と拡張性を示す。

提案手法

• Julia コードから勾配を生成するソースツーソース自動微分アプローチ（Zygote）を用いる。
• バックプロパゲーションのために、ヤコビ行列-ベクトル積で関数出力を包む一般的な微分演算子 J を定義する。
• 複数ディスパッチによるカスタム勾配を可能にする、ユーザー拡張可能な部分微分インタフェースを提供。
• パッケージの素基礎的な変更を要することなく、混合言語および複雑なデータ型を介して微分可能であることを示す。

実験結果

リサーチクエスチョン

• RQ1 制御フロー、変異、複雑なデータ型を持つ一般的な Julia プログラムを微分可能プログラミングシステムが微分できるか？
• RQ2 微分可能プログラミングアプローチは、Julia の既存の ML および科学計算エコシステムとどれだけうまく統合できるか？
• RQ3 CPU・GPU・TPU 上で大規模かつ現実世界の Julia プログラムを微分する場合の性能影響は？
• RQ4 微分可能プログラミングは微分可能なレイトレーシング、量子機械学習、ニューラルSDEs などの応用をエコシステムの組み合わせ性を保ちながら実現できるか？

主な発見

層	総オーバーヘッド	演算数	オーバーヘッド/演算
1	147.0 us	255	576.3 ns
2	280.5 us	491	571.3 ns
3	406.1 us	727	558.6 ns

• Zygote は、フォワードパスとバックワードパスの情報再利用を伴うソースツーソース変換による低オーバーヘッドの自動微分を実現する。
• 隣接演算ごとのオーバーヘッドは、積み重ねられた LSTM の場合平均約 568.8 ns で、PyTorch のオーバーヘッドと競合する。
• 多様なドメイン（レイトレーシング、量子シミュレーション、ニューラルSDE、ファイナンス）を ∂P を用いて微分することは実現可能で有益である。
• 新しい型に対しては、基盤となる AD システムを変更せずに部分微分インタフェースを通じてカスタム勾配を容易に定義できる。
• このアプローチは CPU・GPU・TPU でのトレーニングをサポートし、深層学習タスクにおいて既存の ML フレームワークの性能に匹敵する。