Stabilizing Spiking Neuron Training

TL;DR

The paper tackles the challenge of training stability in spiking neural networks (SNNs) without sacrificing sparsity-driven energy efficiency. It introduces a stability-based framework to guide the initialization of Leaky Integrate-and-Fire (LIF) neurons and to shape surrogate gradients (SGs) before training, aiming to balance forward and backward stability across time. Four concrete conditions (I–IV) link mean and variance of activations, as well as gradient maxima and variance, to SG selection and initial weights, enabling pre-training guidance. Empirically, it shows that higher initial firing rates can improve generalization in deeper networks when paired with a sparsity-encouraging loss, and that the proposed stability constraints improve final accuracy across tasks (SHD, sl-MNIST) and SG shapes, with the derivative of the fast-sigmoid SG often performing robustly. The framework generalizes to other neuron models (ALIF, sLSTM) and offers a principled approach to SG tuning that reduces grid-search demands while promoting energy-efficient, high-performing neuromorphic systems.

Abstract

Stability arguments are often used to prevent learning algorithms from having ever increasing activity and weights that hinder generalization. However, stability conditions can clash with the sparsity required to augment the energy efficiency of spiking neurons. Nonetheless it can also provide solutions. In fact, spiking Neuromorphic Computing uses binary activity to improve Artificial Intelligence energy efficiency. However, its non-smoothness requires approximate gradients, known as Surrogate Gradients (SG), to close the performance gap with Deep Learning. Several SG have been proposed in the literature, but it remains unclear how to determine the best SG for a given task and network. Thus, we aim at theoretically define the best SG, through stability arguments, to reduce the need for grid search. In fact, we show that more complex tasks and networks need more careful choice of SG, even if overall the derivative of the fast sigmoid tends to outperform the other, for a wide range of learning rates. We therefore design a stability based theoretical method to choose initialization and SG shape before training on the most common spiking neuron, the Leaky Integrate and Fire (LIF). Since our stability method suggests the use of high firing rates at initialization, which is non-standard in the neuromorphic literature, we show that high initial firing rates, combined with a sparsity encouraging loss term introduced gradually, can lead to better generalization, depending on the SG shape. Our stability based theoretical solution, finds a SG and initialization that experimentally result in improved accuracy. We show how it can be used to reduce the need of extensive grid-search of dampening, sharpness and tail-fatness of the SG. We also show that our stability concepts can be extended to be applicable on different LIF variants, such as DECOLLE and fluctuations-driven initializations.

Figures (7)

• Figure 1: Surrogate Gradient shapes. To stabilize a network we have to stabilize also the backward pass. However the LIF, as a spiking neuron, has an undefined backward pass and we need Surrogate Gradients (SG) to approximate it. Panel (a) shows the SG investigated in this work, and (b) the tail dependence of our $q$-PseudoSpike SG for $q\in[1.01, 16.85]$. The SG considered are symmetrical around $v_t=y_t-\vartheta=0$, so we only plot half the curve (centered voltage $v_t>0$).
• Figure 2: The choice of SG becomes increasingly important as task and network complexity increase. In order to clearly showcase the problem addressed by our work, and to understand the difficulties brought by SG training, we want to see the impact of training with different SG shapes, and how task and network complexity affect it. This will stress the need for clever rules to apply at initialization to prevent worst case scenarios. Tasks and networks are presented from left to right in order of increasing complexity, where number of classes is used as a proxy of task complexity, and the number of operations is used to quantify network complexity. We perform a grid search over SG shapes, learning rates, tasks and networks. We report lowest validation perplexity after converged training, where perplexity is a loss, so, the lower the better. Panels a-f) show perplexity (y-axis), against learning rate (x-axis). In a-c) we fix the LIF network and change task, while in d-f) we fix the SHD task and change network. Plots b) and d) are repeated for clarity. Panels g-h) show SG sensitivity (y-axis) against task and neural model (x-axis), where we defined sensitivity in Sec. \ref{['sec:sensitivity']}, and it is essentially the variance of the perplexity, across SG shapes and learning rates. a-f) Our results demonstrate that even if different SG shapes tend to agree on the optimal learning rate, the final performance can vary substantially, depending on the SG selection. Specifically, the $\partial$ fast-sigmoid seems the most resilient to changes in the learning rate, as shown in zenke2021remarkable. g-h) Moreover, we observe that the more complex the task or the network, the higher the performance variability we see across SG shapes and learning rates. Our stability criterion provides a method to carefully select SG shapes at initialization and address this issue, promoting better performance and generalization.
• Figure 3: High initialization firing rates can improve generalization with low test firing rates. Our initialization method suggests to set a high firing rate at the beginning of training, which is uncommon in the neuromorphic literature. We study if it is possible to reconcile high initialization firing rates with low firing rates on the test set. We use the SHD task and the $\partial$ fast-sigmoid SG, and measure the correlation $r_l$ of each layer $l$ firing rate with perplexity after training, on the test set. Bold correlation means $p$-value $\leq 0.05$. On the $y$-axis we report perplexity after training on the test set, and on the $x$-axis we report initialization firing rate $\rho_i$, or final firing rate $\rho_f$, meaning the firing rate after training, also evaluated on the test set. On the two left panels, learning starts from different $\rho_i$ without a Sparsity Encouraging Loss Term (SELT), while on the two right panels a target sparsity is encouraged. In both cases, the initial firing rate correlates with final performance, and a low $\rho_f$ is achieved successfully using a SELT. Notice as well that the combination of high initial firing rate and sparsity encouragement resulted in better test loss than on the two panels on the left, suggesting that both factors acted synergistically as a regularization mechanism. We conclude that high initialization firing rates are not necessarily at odds with having sparse activity after training.
• Figure 4: Our stability-focused constraints on the LIF weights and SG shape improve final performance. This figure illustrates our novel method for selecting SG in a Leaky Integrate-and-Fire (LIF) network to improve its stability and performance. We design 4 conditions to stabilize forward and backward pass of a LIF network. (I) requires voltages that promote higher SG values, (II) balances input and recurrent contribution to the voltage, while (III) and (IV) constrain gradient maxima and variance over time. We demonstrate the effectiveness of our method for the LIF network on the SHD task, with an exponential SG and Glorot Uniform initialization, and compare training under the four stability conditions with the baseline without any conditions (shown in gray). Lower and upper panels show validation and test accuracies. Our results show that while condition (II) has the most significant impact on its own, all four conditions combined lead to the best performance. These findings suggest that our theory of LIF stabilit can reduce the need for extensive hyper-parameter search and improve the experimental performance of LIF networks.
• Figure 5: Our stability-based theory predicts optimal SG features on the LIF network. We compare how the features of SG shape predicted by our method stand up against other experimental choices. We conduct the analysis on the LIF network for the sl-MNIST task. Panel (a) shows the performance for different dampening values while setting the sharpness to 1, and vice versa in panel (b). The dashed vertical lines show our theoretical predictions for the exponential SG, (III) for the dampening ($\gamma=0.20\pm 0.02$) and (IV) for the sharpness ($\beta=1.02\pm 0.17$), which agree with the experiments. Dampenings lower than 1 improve performance while the pattern is the opposite for sharpness. Panel (c) shows the performance for different tail-fatness values on the $q$-PseudoSpike SG with $\beta=\gamma=1$. The theoretical prediction gives a close to optimal $q=1.898\pm 0.002$, whereas the best experimental result is $q=1.56$. These findings suggest that our stability-based method predicts good SG features before training, thereby reducing the need for time-consuming hyper-parameter search.

Theorems & Definitions (8)

• Lemma 1
• proof
• Lemma 2
• proof
• Lemma 3
• proof
• Lemma 4
• proof