Slax: A Composable JAX Library for Rapid and Flexible Prototyping of Spiking Neural Networks

• URL: http://arxiv.org/abs/2404.05807v1
• Date: Mon, 8 Apr 2024 18:15:13 GMT
• Title: Slax: A Composable JAX Library for Rapid and Flexible Prototyping of Spiking Neural Networks
• Authors: Thomas M. Summe, Siddharth Joshi,
• Abstract summary: We introduce Slax, a JAX-based library designed to accelerate SNN algorithm design. Slax provides optimized implementations of diverse training algorithms, allowing direct performance comparison.
• Score: 0.19427883580687189
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Recent advances to algorithms for training spiking neural networks (SNNs) often leverage their unique dynamics. While backpropagation through time (BPTT) with surrogate gradients dominate the field, a rich landscape of alternatives can situate algorithms across various points in the performance, bio-plausibility, and complexity landscape. Evaluating and comparing algorithms is currently a cumbersome and error-prone process, requiring them to be repeatedly re-implemented. We introduce Slax, a JAX-based library designed to accelerate SNN algorithm design, compatible with the broader JAX and Flax ecosystem. Slax provides optimized implementations of diverse training algorithms, allowing direct performance comparison. Its toolkit includes methods to visualize and debug algorithms through loss landscapes, gradient similarities, and other metrics of model behavior during training.

Related papers

• Training Artificial Neural Networks by Coordinate Search Algorithm [0.20971479389679332]
  We propose an efficient version of the gradient-free Coordinate Search (CS) algorithm for training neural networks. The proposed algorithm can be used with non-differentiable activation functions and tailored to multi-objective/multi-loss problems. Finding the optimal values for weights of ANNs is a large-scale optimization problem.
  arXiv  Detail & Related papers  (2024-02-20T01:47:25Z)
• Improving Stability and Performance of Spiking Neural Networks through Enhancing Temporal Consistency [9.545711665562715]
  Spiking neural networks have gained significant attention due to their brain-like information processing capabilities. Current training algorithms tend to overlook the differences in output distribution at various timesteps. We have designed a method to enhance the temporal consistency of outputs at different timesteps.
  arXiv  Detail & Related papers  (2023-05-23T15:50:07Z)
• The Cascaded Forward Algorithm for Neural Network Training [61.06444586991505]
  We propose a new learning framework for neural networks, namely Cascaded Forward (CaFo) algorithm, which does not rely on BP optimization as that in FF. Unlike FF, our framework directly outputs label distributions at each cascaded block, which does not require generation of additional negative samples. In our framework each block can be trained independently, so it can be easily deployed into parallel acceleration systems.
  arXiv  Detail & Related papers  (2023-03-17T02:01:11Z)
• Intelligence Processing Units Accelerate Neuromorphic Learning [52.952192990802345]
  Spiking neural networks (SNNs) have achieved orders of magnitude improvement in terms of energy consumption and latency. We present an IPU-optimized release of our custom SNN Python package, snnTorch.
  arXiv  Detail & Related papers  (2022-11-19T15:44:08Z)
• Improved Algorithms for Neural Active Learning [74.89097665112621]
  We improve the theoretical and empirical performance of neural-network(NN)-based active learning algorithms for the non-parametric streaming setting. We introduce two regret metrics by minimizing the population loss that are more suitable in active learning than the one used in state-of-the-art (SOTA) related work.
  arXiv  Detail & Related papers  (2022-10-02T05:03:38Z)
• Tricks and Plugins to GBM on Images and Sequences [18.939336393665553]
  We propose a new algorithm for boosting Deep Convolutional Neural Networks (BoostCNN) to combine the merits of dynamic feature selection and BoostCNN. We also propose a set of algorithms to incorporate boosting weights into a deep learning architecture based on a least squares objective function. Experiments show that the proposed methods outperform benchmarks on several fine-grained classification tasks.
  arXiv  Detail & Related papers  (2022-03-01T21:59:00Z)
• A Continuous Optimisation Benchmark Suite from Neural Network Regression [0.0]
  Training neural networks is an optimisation task that has gained prominence with the recent successes of deep learning. gradient descent variants are by far the most common choice with their trusted good performance on large-scale machine learning tasks. We contribute CORNN, a suite for benchmarking the performance of any continuous black-box algorithm on neural network training problems.
  arXiv  Detail & Related papers  (2021-09-12T20:24:11Z)
• Gone Fishing: Neural Active Learning with Fisher Embeddings [55.08537975896764]
  There is an increasing need for active learning algorithms that are compatible with deep neural networks. This article introduces BAIT, a practical representation of tractable, and high-performing active learning algorithm for neural networks.
  arXiv  Detail & Related papers  (2021-06-17T17:26:31Z)
• Evolving Reinforcement Learning Algorithms [186.62294652057062]
  We propose a method for meta-learning reinforcement learning algorithms. The learned algorithms are domain-agnostic and can generalize to new environments not seen during training. We highlight two learned algorithms which obtain good generalization performance over other classical control tasks, gridworld type tasks, and Atari games.
  arXiv  Detail & Related papers  (2021-01-08T18:55:07Z)
• Communication-Efficient Distributed Stochastic AUC Maximization with Deep Neural Networks [50.42141893913188]
  We study a distributed variable for large-scale AUC for a neural network as with a deep neural network. Our model requires a much less number of communication rounds and still a number of communication rounds in theory. Our experiments on several datasets show the effectiveness of our theory and also confirm our theory.
  arXiv  Detail & Related papers  (2020-05-05T18:08:23Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.