Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

ADMM-Based Training for Spiking Neural Networks

Giovanni Perin, Cesare Bidini, Riccardo Mazzieri, Michele Rossi

TL;DR

This work addresses the training bottleneck of sensing-and-spiking neural networks (SNNs) caused by the non-differentiable spike activation and surrogate-gradient limitations. It introduces an ADMM-based, gradient-free optimization framework that treats SNN training as a constrained, model-based problem with variables for weights, membrane potentials, and spikes, plus a dedicated subroutine to manage the Heaviside nonlinearity and exact last-layer dynamics. Key contributions include (i) a new ADMM optimizer tailored for SNNs, (ii) problem relaxation with closed-form subproblem updates, (iii) a subroutine to optimally handle $H_ heta$, and (iv) a proof-of-concept demonstration showing convergence and potential, with discussion of extensions to deeper and varied layer types. The approach offers a scalable alternative to surrogate-gradient methods, enables parallelization and potential federated learning, and lays out concrete paths to scale the method to larger SNNs and different architectures.

Abstract

In recent years, spiking neural networks (SNNs) have gained momentum due to their high potential in time-series processing combined with minimal energy consumption. However, they still lack a dedicated and efficient training algorithm. The popular backpropagation with surrogate gradients, adapted from stochastic gradient descent (SGD)-derived algorithms, has several drawbacks when used as an optimizer for SNNs. Specifically, the approximation introduced by the use of surrogate gradients leads to numerical imprecision, poor tracking of SNN firing times at training time, and, in turn, poor scalability. In this paper, we propose a novel SNN training method based on the alternating direction method of multipliers (ADMM). Our ADMM-based training aims to solve the problem of the SNN step function's non-differentiability by taking an entirely new approach with respect to gradient backpropagation. For the first time, we formulate the SNN training problem as an ADMM-based iterative optimization, derive closed-form updates, and empirically show the optimizer's convergence, its great potential, and discuss future and promising research directions to improve the method to different layer types and deeper architectures.

ADMM-Based Training for Spiking Neural Networks

TL;DR

This work addresses the training bottleneck of sensing-and-spiking neural networks (SNNs) caused by the non-differentiable spike activation and surrogate-gradient limitations. It introduces an ADMM-based, gradient-free optimization framework that treats SNN training as a constrained, model-based problem with variables for weights, membrane potentials, and spikes, plus a dedicated subroutine to manage the Heaviside nonlinearity and exact last-layer dynamics. Key contributions include (i) a new ADMM optimizer tailored for SNNs, (ii) problem relaxation with closed-form subproblem updates, (iii) a subroutine to optimally handle , and (iv) a proof-of-concept demonstration showing convergence and potential, with discussion of extensions to deeper and varied layer types. The approach offers a scalable alternative to surrogate-gradient methods, enables parallelization and potential federated learning, and lays out concrete paths to scale the method to larger SNNs and different architectures.

Abstract

In recent years, spiking neural networks (SNNs) have gained momentum due to their high potential in time-series processing combined with minimal energy consumption. However, they still lack a dedicated and efficient training algorithm. The popular backpropagation with surrogate gradients, adapted from stochastic gradient descent (SGD)-derived algorithms, has several drawbacks when used as an optimizer for SNNs. Specifically, the approximation introduced by the use of surrogate gradients leads to numerical imprecision, poor tracking of SNN firing times at training time, and, in turn, poor scalability. In this paper, we propose a novel SNN training method based on the alternating direction method of multipliers (ADMM). Our ADMM-based training aims to solve the problem of the SNN step function's non-differentiability by taking an entirely new approach with respect to gradient backpropagation. For the first time, we formulate the SNN training problem as an ADMM-based iterative optimization, derive closed-form updates, and empirically show the optimizer's convergence, its great potential, and discuss future and promising research directions to improve the method to different layer types and deeper architectures.
Paper Structure (14 sections, 10 equations, 4 figures, 1 table, 2 algorithms)

This paper contains 14 sections, 10 equations, 4 figures, 1 table, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. SNNs and training algorithms
  4. Gradient-free training with the ADMM
  5. Optimization problem formulation
  6. Problem relaxation
  7. ADMM-based solution
  8. Weights update
  9. Pre-activations update
  10. Post-activations update
  11. Training algorithm
  12. Numerical simulations
  13. Discussion
  14. Concluding remarks

Figures (4)

  • Figure 1: Graphical representation of the SNN model used for the problem formulation in this paper.
  • Figure 2: Train accuracy obtained in 1,000 ADMM iterations varying the number of hidden layers.
  • Figure 3: Objective functions of the ADMM optimization: relaxed Lagrangian (left) and target loss function (right). Training of the SNN with two hidden layers.
  • Figure 4: Constraint residuals normalized by their dimensions ($\sqrt{T \times M \times n_{l} }$). Exact constraint primal residual of the last membrane potential (left), norm of the soft constraints relative to the membrane potential dynamics per layer (middle), and norm of the soft constraints relative to the neurons' spikes per layer (right). Training of the SNN with two hidden layers.