EXODUS: Stable and Efficient Training of Spiking Neural Networks
Spiking Neural Networks (SNNs) are gaining significant traction in machine
learning tasks where energy-efficiency is of utmost importance. Training such
networks using the state-of-the-art back-propagation through time (BPTT) is,
however, very time-consuming. Previous work by Shrestha and Orchard [2018]
employs an efficient GPU-accelerated back-propagation algorithm called SLAYER,
which speeds up training considerably. SLAYER, however, does not take into
account the neuron reset mechanism while computing the gradients, which we
argue to be the source of numerical instability. To counteract this, SLAYER
introduces a gradient scale hyperparameter across layers, which needs manual
tuning. In this paper, (i) we modify SLAYER and design an algorithm called
EXODUS, that accounts for the neuron reset mechanism and applies the Implicit
Function Theorem (IFT) to calculate the correct gradients (equivalent to those
computed by BPTT), (ii) we eliminate the need for ad-hoc scaling of gradients,
thus, reducing the training complexity tremendously, (iii) we demonstrate, via
computer simulations, that EXODUS is numerically stable and achieves a
comparable or better performance than SLAYER especially in various tasks with
SNNs that rely on temporal features. Our code is available at
https://github.com/synsense/sinabs-exodus.