Spyx：用于脉冲神经网络即时编译优化的库
神经与进化计算
2024-03-01 v1 机器学习
摘要
随着人工智能在现代社会中的作用日益重要,深度神经网络的高效训练与部署已成为关键关注领域。近期注意力驱动的大型神经网络架构的发展,推动了AI加速器的开发,以便训练大型多数十亿参数模型。尽管这些强大网络常在生产环境中造成高执行成本。仿生计算提供了有前景的替代方案。通过利用时序稀疏计算,脉冲神经网络(SNNs)可通过缩减和低功耗硬件资源来提高能源效率。然而,由于其时滞性质,SNN的训练可能具有挑战性,难以充分利用现代AI加速器的大规模并行性。为促进SNN架构和动态研究,研究人员寻求将基于Python的深度学习框架(如PyTorch或TensorFlow)与自定义实现的计算内核桥接。本文介绍Spyx,一个全新的轻量级SNN仿真与优化库,采用JAX开发。通过在现代加速器的广阔vRAM中预置数据,并采用大量即时编译,Spyx允许SNN优化在NVIDIA GPU或Google TPU上作为统一的低级程序执行,从而实现最佳硬件利用率,优于许多现有SNN训练框架,同时保持相当大的灵活性。
引用
@article{arxiv.2402.18994,
title = {Spyx: A Library for Just-In-Time Compiled Optimization of Spiking Neural Networks},
author = {Kade M. Heckel and Thomas Nowotny},
journal= {arXiv preprint arXiv:2402.18994},
year = {2024}
}