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

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

Parallel Computing and Optimization Techniques参考文献 39被引用数 63

ひとこと要約

Zygoteは、Juliaで実装された微分可能プログラミングシステムであり、制御構造、再帰、可変性を含む一般プログラム構造に対して、ユーザーの介入なしに自動微分を可能にする。高性能で段階的な計算をサポートし、科学計算ライブラリとシームレスに統合され、表現力があり微分可能なプログラミングを通じて機械学習と科学計算を統合するフレームワークを提供する。

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.

研究の動機と目的

• 一般プログラム構造にわたる微分可能プログラミングを可能にすることで、機械学習と科学計算を統合すること。
• 最小限のユーザー負荷でJuliaにおける自動微分をサポートし、再構築や段階的計算の必要性を回避すること。
• 統一的で表現力のあるプログラミングモデルを通じて、科学計算ライブラリを機械学習モデルにシームレスに統合すること。
• 混合モード、複素数、チェックポイント付き微分などの高度な技術への自動微分の拡張を実現すること。
• 微分可能プログラミングが、実世界の科学的および機械学習ワークロードをサポートする性能と汎用性を両立できることを示すこと。

提案手法

• Zygoteは、Juliaのメタディスpatchとメタプログラミング機能を活用して、効率的な勾配計算を生成することで、高性能コードをコンパイルする。
• 任意のJuliaプログラム（制御構造、再帰、可変状態を含む）を自動的に微分するため、ソースツーソース変換を用いる。
• システムは、抽象構文木の記号的変換を通じて、混合モードおよび複素数微分をサポートする逆モード自動微分エンジンを採用している。
• チェックポイント処理を実装し、中間値を適切に再計算することで、バックプロパゲーション中のメモリ使用量を削減する。
• Zygoteは、Juliaの既存の数値計算および科学計算ライブラリエコシステムと統合されており、複雑な科学的ワークロードのエンドツーエンド微分可能性を実現する。
• フレームワークは拡張性を備えており、最小限のボイラープレートでユーザーがカスタム微分可能演算を定義できる。

実験結果

リサーチクエスチョン

• RQ1 微分可能プログラミングシステムは、コードの再構築を要せず、制御構造、再帰、可変性といった一般プログラミング構造をサポートできるか？
• RQ2 微分可能プログラミングシステムは、複雑な科学計算ワークロードの勾配を、どの程度効率的にコンパイル・実行できるか？
• RQ3 一般プログラミング言語において、混合モードおよび複素数微分の自動微分をどの程度拡張できるか？
• RQ4 微分可能プログラミングは、科学計算ライブラリを機械学習モデルにシームレスに統合するために使用できるか？
• RQ5 高水準言語（Julia）で全プログラム微分可能にすることにより、どのようなパフォーマンスと使いやすさのトレードオフが生じるか？

主な発見

• Zygoteは、複雑な制御構造や可変性を含む一般Juliaプログラムを、ユーザーの介入やコード再構築なしに正常に微分する。
• システムは、効率的で低レベルのコードをコンパイルすることで、ほとんどの場合でランタイムのオーバーヘッドを回避し、高いパフォーマンスを達成する。
• 混合モードおよび複素数微分は、プログラムの抽象構文木の記号的変換を通じてネイティブにサポートされている。
• チェックポイント処理は効果的に統合されており、バックプロパゲーション中のメモリ消費量を削減しながら、計算効率を維持している。
• フレームワークにより、常微分方程式ソルバーや線形代数ライブラリなどの科学計算ライブラリを、直接微分可能モデル内で使用可能にし、表現力と再利用性を向上させている。
• Zygoteは、微分可能プログラミングが、1つのパフォーマンスと使いやすさに優れたシステムとして、機械学習と科学計算を統合できることを示している。