[論文レビュー] Spyx: A Library for Just-In-Time Compiled Optimization of Spiking Neural Networks

Kade Heckel, Thomas Nowotny|arXiv (Cornell University)|Feb 29, 2024

Advanced Memory and Neural Computing被引用数 5

ひとこと要約

Spyx は、JAX ベースの軽量な SNN ライブラリで、JIT コンパイルと vRAM データステージングを利用して、GPU/TPU 上でスパイキングニューラルネットワークを効率的に訓練・シミュレーションします。低レベルの CUDA カーネルに匹敵しつつ、柔軟性を保つことを目指します。

ABSTRACT

As the role of artificial intelligence becomes increasingly pivotal in modern society, the efficient training and deployment of deep neural networks have emerged as critical areas of focus. Recent advancements in attention-based large neural architectures have spurred the development of AI accelerators, facilitating the training of extensive, multi-billion parameter models. Despite their effectiveness, these powerful networks often incur high execution costs in production environments. Neuromorphic computing, inspired by biological neural processes, offers a promising alternative. By utilizing temporally-sparse computations, Spiking Neural Networks (SNNs) offer to enhance energy efficiency through a reduced and low-power hardware footprint. However, the training of SNNs can be challenging due to their recurrent nature which cannot as easily leverage the massive parallelism of modern AI accelerators. To facilitate the investigation of SNN architectures and dynamics researchers have sought to bridge Python-based deep learning frameworks such as PyTorch or TensorFlow with custom-implemented compute kernels. This paper introduces Spyx, a new and lightweight SNN simulation and optimization library designed in JAX. By pre-staging data in the expansive vRAM of contemporary accelerators and employing extensive JIT compilation, Spyx allows for SNN optimization to be executed as a unified, low-level program on NVIDIA GPUs or Google TPUs. This approach achieves optimal hardware utilization, surpassing the performance of many existing SNN training frameworks while maintaining considerable flexibility.

研究の動機と目的

• 省エネルギーなニューロモorphic コンピューティングのための Spiking Neural Networks (SNNs) の効率的な訓練とデプロイを促進する。
• JAX ベースで PyTorch に優しい API を提供し、JIT コンパイルと Haiku 連携を通じて SNN 研究を加速する。
• 最小限の低レベルカーネルプログラミングで、柔軟なニューロンモデルと surrogate gradients を実現する。

提案手法

• JAX/Haiku の上に Spyx を設計し、JIT の機会を最大化し、機能的で状態を持たないワークフローを維持する。
• 高階関数を用いて surrogate gradient 関数を実装し、カスタム forward/backward の定義を可能にする。
• ニューロンモデル（例：Leaky-Integrate-and-Fire）を Haiku RNN コアとして、JAX ベースのダイナミクスと動的/静的展開を備える。
• オン-GPU データパック/アンパック、動的解凍、GPU 中心のデータ拡張/シャッフルを含む、データ取り扱いの最適化を提供する。
• Neuromorphic Intermediate Representation (NIR) との統合により、他のフレームワークやハードウェアとの間で SNN の容易なインポート/エクスポートを実現する。

実験結果

リサーチクエスチョン

• RQ1 ベースの CUDA カーネルを特注しなくても、JAX ベースの SNN フレームワークは PyTorch ベースの SNN ライブラリと同等またはそれを上回る訓練性能を提供できるか。
• RQ2 surrogate gradient の選択とニューロンモデルの実装が、JAX/JIT コンパイル設定で訓練速度と精度にどう影響するか。
• RQ3 データ処理とコンパイル戦略（例：データパック/アンパック、静的 vs. 動的展開）により、SNN 訓練でハードウェア利用率を最大化するには何を採用すべきか？
• RQ4 Spyx で訓練された SNN モデルを NIR に従ってシリアライズ/デシリアライズして、ニューロモorphic ハードウェアへ展開するにはどの程度可能か？

主な発見

• Spyx は、JAX JIT コンパイルと Haiku ベースの RNN コアを活用して、カスタム CUDA カーネルを必要とせず、著名な SNN フレームワークに対して競争力のある訓練速度を達成します。
• このライブラリは、柔軟な surrogate gradient とニューロンモデルを、コンパクトでモジュール化された API でサポートし、迅速な実験を可能にします。
• SHD および NMNIST のベンチマークでは、Spyx は複数の PyTorch ベースのライブラリよりも大幅なスピードアップを示し、ahead-of-time コンパイル後の NMNIST では SpikingJelly の性能の 5% 内に収まります。
• データ処理の最適化（例：オン-GPU のパック/アンパック、時系列データ圧縮）は I/O レイテンシとメモリ使用量を低減し、スループットの向上に寄与します。
• NIR ベースのエクスポート/インポートにより、ソフトウェアとハードウェアのターゲット間でモデルの移植性が高まり、ニューロモorphic プラットフォームへのデプロイを支援します。
• 設計は最小限の CPU 作業と最大限の GPU 利用を強調し、純粋な JAX ベースのエコシステム内で CUDA カーネルレベルに近いパフォーマンスを達成します。