[论文解读] 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.

研究动机与目标

• 推动对 Spiking Neural Networks (SNNs) 的高效训练与部署，以实现能源高效的类脑计算。
• 提供一个基于 JAX、兼容 PyTorch 的 API，通过 JIT 编译和 Haiku 集成来加速 SNN 研究。
• 在尽量减少底层内核编程的前提下，实现灵活的神经元模型和代理梯度。

提出的方法

• 在 JAX/Haiku 之上设计 Spyx，以最大化 JIT 机会并保持函数式、无状态的工作流。
• 通过高阶函数实现代理梯度函数，以允许自定义前向/后向定义。
• 以 Haiku RNN core 的形式提供神经元模型（例如 Leaky-Integrate-and-Fire），具有基于 JAX 的动力学和动态/静态展开。
• 提供数据处理优化，包括在 GPU 上进行数据打包/解包、动态解压缩，以及以 GPU 为中心的增强/乱序。
• 与 Neuromorphic Intermediate Representation (NIR) 集成，便于 SNNs 与其他框架及硬件之间的导入/导出。

实验结果

研究问题

• RQ1 在没有定制 CUDA 内核的情况下，基于 JAX 的 SNN 框架能否达到与 PyTorch 基于的 SNN 库相当或更高的训练性能？
• RQ2 在 JAX/JIT 编译环境中，代理梯度的选择和神经元模型实现如何影响训练速度和准确性？
• RQ3 哪些数据处理与编译策略（例如数据打包、静态展开 vs 动态展开）能最大化 SNN 训练的硬件利用率？
• RQ4 在多大程度上，Spyx 训练的 SNN 模型可以序列化为/从 Neuromorphic Intermediate Representation (NIR) 导出，以在类脑硬件上部署？

主要发现

• Spyx 在训练速度方面与知名 SNN 框架具有竞争力，利用 JAX JIT 编译和 Haiku 基于 RNN 的核心，而无需定制 CUDA 内核。
• 该库支持灵活的代理梯度和神经元模型，具备紧凑、模块化的 API，便于快速实验。
• 在 SHD 和 NMNIST 基准测试中，Spyx 相较于若干基于 PyTorch 的库显示出显著速度提升，并且在 NMNIST 的 ahead-of-time 编译后性能仅落后 SpikingJelly 约 5%。
• 数据处理优化（如在 GPU 上打包/解包、时间序列数据压缩）降低 I/O 延迟和内存使用，从而提高吞吐量。
• 基于 NIR 的导出/导入实现了跨软件和硬件目标的模型可移植性，便于部署到类脑平台。
• 设计强调最小化 CPU 工作量、最大化 GPU 利用率，在纯粹的 JAX 生态系统中实现接近 CUDA 内核级的性能。