SNNAX — Spiking Neural Networks in JAX

要約

スパイキング ニューラル ネットワーク (SNN) シミュレーターは、生物学にインスピレーションを得たモデルやニューロモーフィック ハードウェア アーキテクチャのプロトタイプを作成し、そのパフォーマンスを予測するために不可欠なツールです。
このようなツールでは、使いやすさと柔軟性が重要ですが、特に SNN のシミュレーションに固有の複雑さを考慮すると、シミュレーション速度も重要です。
ここでは、PyTorch のような直観性と JAX のような実行速度でそのようなモデルをシミュレーションおよびトレーニングするための JAX ベースのフレームワークである SNNAX を紹介します。
SNNAX モデルは、目的のモデル仕様やターゲットのニューロモーフィック ハードウェアに合わせて簡単に拡張およびカスタマイズできます。
さらに、SNNAX は、柔軟な自動微分やジャストインタイム コンパイルなど、SNN のトレーニングと展開を最適化するための重要な機能を提供します。
SNNAX を、SNN のプログラミングに使用される他の一般的に使用される機械学習 (ML) フレームワークと評価および比較します。
SNNAX で SNN をシミュレートするための主要なパフォーマンス メトリクス、ベスト プラクティス、文書化された例を提供し、文献で使用されているいくつかのベンチマークを実装します。

要約(オリジナル)

Spiking Neural Networks (SNNs) simulators are essential tools to prototype biologically inspired models and neuromorphic hardware architectures and predict their performance. For such a tool, ease of use and flexibility are critical, but so is simulation speed especially given the complexity inherent to simulating SNN. Here, we present SNNAX, a JAX-based framework for simulating and training such models with PyTorch-like intuitiveness and JAX-like execution speed. SNNAX models are easily extended and customized to fit the desired model specifications and target neuromorphic hardware. Additionally, SNNAX offers key features for optimizing the training and deployment of SNNs such as flexible automatic differentiation and just-in-time compilation. We evaluate and compare SNNAX to other commonly used machine learning (ML) frameworks used for programming SNNs. We provide key performance metrics, best practices, documented examples for simulating SNNs in SNNAX, and implement several benchmarks used in the literature.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google