[論文レビュー] SymJAX: symbolic CPU/GPU/TPU programming

Randall Balestriero|arXiv (Cornell University)|May 21, 2020

Music and Audio Processing参考文献 5被引用数 4

ひとこと要約

SymJAX は、JAX と XLA を基盤に構築された記号的プログラミングフレームワークであり、カーネルの統合、メモリ操作の削減、自動微分のサポートにより、CPU、GPU、TPU における効率的で最適化された計算グラフを実現する。JIT コンパイルと広範なハードウェアサポートを備えた、PyTorch に似た高水準なAPIを提供し、グラフ最適化と記号的計算によってパフォーマンス向上を実証する。

ABSTRACT

SymJAX is a symbolic programming version of JAX simplifying graph input/output/updates and providing additional functionalities for general machine learning and deep learning applications. From an user perspective SymJAX provides a la Theano experience with fast graph optimization/compilation and broad hardware support, along with Lasagne-like deep learning functionalities.

研究の動機と目的

• 計算グラフの構築、入出力の取り扱い、パラメータの更新を簡素化する記号的プログラミングインターフェースをディープラーニングに提供すること。
• 複数のカーネルを単一の操作に統合することで、計算グラフを最適化し、メモリ転送を削減してパフォーマンスを向上させること。
• XLA を通じて、CPU、GPU、TPU ハードウェアにおいて、効率的な JIT コンパイルと自動微分をサポートすること。
• JAX のパフォーマンスとコンパイルの利点を保ちつつ、PyTorch に似た直感的で高水準な API を提供すること。
• 最適化と微分を完全にサポートする、オプティマイザ、損失関数、レイヤー部品を内蔵したエンドツーエンドのトレーニングおよび推論パイプラインを可能にすること。

提案手法

• JAX の XLA バックエンドを介して記号的テンソルを使用し、静的グラフの構築と最適化を可能にする。
• XLA を通じたカーネル統合により、演算を統合し、中間メモリの割り当てを最小限に抑え、実行速度を向上させる。
• JAX の XLA コンパイラを介した JIT コンパイルを活用し、計算グラフの高速化を実現する。
• SymJAX の勾配関数を通じて自動微分をサポートし、記号的式に対する誤差逆伝播を可能にする。
• `sj.function` を通じた関数的 API を提供し、指定された出力と更新を伴うグラフのコンパイルと実行を可能にする。
• `sj.layers` および `sj.optimizers` を通じて、畳み込み層（Conv2D）、バッチノーマライゼーション、ReLU などのモジュラーなニューラルネットワークコンポONENTを提供し、モデル構築を支援する。

実験結果

リサーチクエスチョン

• RQ1 JAX と XLA を基盤に構築した記号的プログラミングインターフェースは、複数のアクセラレータをサポートしながら、PyTorch と同等のパフォーマンスと使いやすさを達成できるか？
• RQ2 カーネル統合は、ディープラーニングワークロードにおけるメモリ操作の削減と、推論およびトレーニング速度の向上にどの程度効果的か？
• RQ3 自動微分を伴う記号的計算を、複雑なディープラーニングモデルに対して使いやすく、合成可能にするにはどの程度の水準に到達できるか？
• RQ4 XLA ベースのコンパイルと JIT コンパイルの統合は、CPU、GPU、TPU の各ハードウェアにおけるディープラーニングモデルのパフォーマンスにどのように寄与するか？
• RQ5 JAX の上に、最小限のパフォーマンスオーバーヘッドで、最適化と微分を完全にサポートする、PyTorch に似た高水準な API を構築できるか？

主な発見

• SymJAX は、多次元正規分布の確率密度関数（PDF）の計算を正常にコンパイル・最適化し、遅延評価が可能で、正しい出力が得られた。
• 確率的勾配降下法の実験では、SymJAX が正しく勾配を計算し、パラメータを更新した。10イテレーションで損失が約 0.00847 から約 0.0065 に低下した。
• CIFAR-10 では、100エポック後、SymJAX でトレーニングされたモデルがテスト損失約 0.815、テスト精度約 72.1% を達成し、最適化された計算による有効なトレーニングを実証した。
• フレームワークは、畳み込み層、バッチノーマライゼーション、プーリング、全結合層を含む複雑なモデルアーキテクチャをサポートし、正しい形状伝播とパラメータ追跡が行われた。
• `sj.function` の使用により、更新を伴うグラフの効率的なコンパイルと実行が可能となり、損失と精度がエポックごとに追跡されるトレーニングループでその有効性が示された。
• カーネル統合と JIT コンパイルにより、メモリ操作が顕著に削減され、パフォーマンスが向上した。CIFAR-10 トレーニングパイプラインの効率的実行がその証左となった。