[논문 리뷰] Spyx: A Library for Just-In-Time Compiled Optimization of Spiking Neural Networks

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

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 연구를 가속화합니다.
• 저수준 커널 프로그래밍을 최소화하면서 유연한 뉴런 모델과 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 JAX 기반 SNN 프레임워크가 bespoke CUDA 커널 없이 PyTorch 기반 SNN 라이브러리와 비교해 동등하거나 더 우수한 학습 성능을 낼 수 있을까요?
• RQ2 surrogate gradient 선택과 뉴런 모델 구현이 JAX/JIT-컴파일 설정에서 학습 속도와 정확도에 어떤 영향을 미치나요?
• RQ3 데이터 패킹, 정적 대 동적 언롤링과 같은 데이터 처리 및 컴파일 전략이 SNN 학습을 위한 하드웨어 활용을 어떻게 극대화하나요?
• RQ4 Spyx에서 학습된 SNN 모델을 Neuromorphic Intermediate Representation (NIR)로 직렬화/역직렬화하여 신경모사 하드웨어에 배치하는 정도는 어느 정도인가요?

주요 결과

• Spyx는 JAX JIT 컴파일과 Haiku 기반 RNN 코어를 활용하여 커스텀 CUDA 커널 없이도 주요 SNN 프레임워크에 비해 경쟁력 있는 학습 속도를 달성합니다.
• 라이브러리는 유연한 surrogate gradients 및 뉴런 모델을 컴팩트하고 모듈식 API로 지원하여 빠른 실험을 가능하게 합니다.
• SHD 및 NMNIST 벤치마크에서 Spyx는 다수의 PyTorch 기반 라이브러리에 비해 상당한 속도 향상을 보이고, ahead-of-time 컴파일 후 NMNIST에서 SpikingJelly 성능의 5% 이내에 들어갑니다.
• 데이터 처리 최적화(예: 온-GPU 패킹/언패킹, 시간적 데이터 압축)가 I/O 지연 및 메모리 사용을 줄여 처리량을 높입니다.
• NIR 기반 수출/수입으로 소프트웨어 및 하드웨어 대상 간 모델 이식성을 가능하게 하여 신경모사 플랫폼으로의 배치를 돕습니다.
• 설계는 CPU 작업을 최소화하고 GPU 활용을 극대화하는 데 중점을 두어 Purely JAX 기반 생태계에서 CUDA 커널 수준의 성능에 근접합니다.