Slax:用于快速灵活地构建脉冲神经网络原型的可组合JAX库
神经与进化计算
2024-04-10 v1
摘要
训练脉冲神经网络(SNN)算法的最新进展通常利用其独特的动力学。尽管使用替代梯度的随时间反向传播(BPTT)主导了该领域,但丰富的替代方案可以将算法定位在性能、生物合理性和复杂性图景中的各个位置。评估和比较算法目前是一个繁琐且易错的过程,需要将它们反复重新实现。我们引入了Slax,这是一个基于JAX的库,旨在加速SNN算法设计,并与更广泛的JAX和Flax生态系统兼容。Slax提供了多种训练算法的优化实现,允许直接进行性能比较。其工具包包括通过损失景观、梯度相似性以及训练期间模型行为的其他指标来可视化和调试算法的方法。
引用
@article{arxiv.2404.05807,
title = {Slax: A Composable JAX Library for Rapid and Flexible Prototyping of Spiking Neural Networks},
author = {Thomas M. Summe and Siddharth Joshi},
journal= {arXiv preprint arXiv:2404.05807},
year = {2024}
}
备注
13 pages, 11 figures, early draft