用JAX实现高速脉冲神经网络模拟，兼具易用与高效。
SNNAX -- Spiking Neural Networks in JAX
- 基于JAX构建，支持自动微分与即时编译提升速度。
- 可灵活扩展，适配不同神经网络模型与类脑硬件。
- 提供训练优化工具和文献基准测试，适合研究者使用。
脉冲神经网络（SNN）模拟器是原型设计生物启发模型和类脑硬件架构的关键工具，其性能预测至关重要。这类工具需兼顾易用性、灵活性与仿真速度，尤其在复杂度较高的情况下。本文提出SNNAX，一个基于JAX的框架，用于模拟与训练SNN，具备类似PyTorch的直观性与类似JAX的执行速度。SNNAX模型可轻松扩展和定制，以匹配特定模型需求及目标类脑硬件。此外，该框架提供灵活的自动微分机制与即时编译功能，助力优化训练与部署。我们对SNNAX与其他常用机器学习框架在编程SNN方面的表现进行了评估与对比，提供了关键性能指标、最佳实践建议、文档示例，并实现了多个文献中常用的基准测试。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。