Pruning of Deep Spiking Neural Networks through Gradient Rewiring

1 Introduction

Spiking Neural Networks (SNNs), as the third generation of Artificial Neural Networks (ANN) [Maass, 1997], have increasingly aroused researchers' great interest in recent years due to their high biological plausibility, temporal information processing capability, and low power consumption on neuromorphic chips [Gerstner and Kistler, 2002; Roy et al., 2019].

∗Corresponding author

As these chips are usually resource-constrained, the compression of SNNs is thus of great importance to lower computation cost to make the deployment on neuromorphic chips more feasible, and realize the practical application of SNNs.

Researchers have made a lot of efforts on the pruning of SNNs. Generally, most methods learn from the existing compression techniques in the field of ANNs and directly apply them to SNNs, such as directly truncating weak weights below the threshold [Neftci et al., 2016; Rathi et al., 2019; Liu et al., 2020] and using Lottery Ticket Hypothesis [Frankle and Carbin, 2019] in SNNs [Martinelli et al., 2020]. Some works use other metrics, such as levels of weight development [Shi et al., 2019] and similarity between spike trains [Wu et al., 2019a], as guidance for removing insignificant or redundant weights. Notwithstanding some spike-based metrics for pruning are introduced in these works, the results on shallow SNNs are not persuasive enough.

There are some attempts to mimic the way the human brain learns to enhance the pruning of SNNs. Qi et al. [2018] enabled structure learning of SNNs by adding connection gates for each synapse. They adopted the biological plausible local learning rules, Spike-Timing-Dependent Plasticity (STDP), to learn gated connection while using Tempotron rules [G¨utig and Sompolinsky, 2006] for regular weight optimization. However, the method only applies to the SNN with a two-layer structure. Deng et al. [2019] exploited the ADMM optimization tool to prune the connections and optimize network weights alternately. Kappel et al. [2015] proposed synaptic sampling to optimize network construction and weights by modeling the spine motility of spiking neurons as Bayesian learning. Bellec et al. went a further step and proposed Deep Rewiring (Deep R) as a pruning algorithm for ANNs [2018a] and Long Short-term Memory Spiking Neural Networks (LSNNs) [2018b], which was then deployed on SpiNNaker 2 prototype chips [Liu et al., 2018]. Deep R is an on-the-fly pruning algorithm, which indicates its ability to reach target connectivity without retraining. Despite these methods attempting to combine learning and pruning in a unified framework, they focus more on pruning, rather than the rewiring process. Actually, the pruned connection is hard to regrow, limiting the diversity of network structures and thus the performance of the pruned SNNs.

In this paper, inspired by synaptogenesis and synapse elimination in the neural system [Huttenlocher and Nelson, 1994;

Petzoldt and Sigrist, 2014], we propose gradient rewiring (Grad R), a joint learning algorithm of connectivity and weight for SNNs, by introducing a new synaptic parameter to unify the representation of synaptic connectivity and synaptic weights. By redefining the gradient to the synaptic parameter, Grad R takes full advantage of the competition between pruning and regrowth of connections to explore the optimal network structure. To show Grad R can optimize synaptic parameters of SNNs, we theoretically prove that the angle between the accurate gradient and redefined gradient is less than 90◦ . We further derive the learning rule and realize the pruning of deep SNNs. We evaluate our algorithm on the MNIST and CIFAR-10 dataset and achieves state-of-the-art performance. On CIFAR-10, Grad R reaches a classification accuracy of 92.03% (merely <1% accuracy loss) when constrained to ∼5% connectivity. Further, it achieves a classification accuracy of 89.32% (∼3.5% accuracy loss) under unprecedented 0.73% connectivity, which reveals remarkable structure refining capability in SNNs. Our work suggests that there exists extremely high redundancy in deep SNNs.

2 Methods

In this section, we present the pruning algorithm for SNNs in detail. We first introduce the spiking neural model and its discrete computation mechanism. Then we propose the network structure of the SNNs used in this work. Finally, we present the gradient rewiring framework to learn the connectivity and weights of SNNs jointly and derive the detailed learning rule.

2.1 Spiking Neural Model

The spiking neuron is the fundamental computing unit of SNNs. Here we consider the Leaky Integrate-and-Fire (LIF) model [Gerstner and Kistler, 2002], which is one of the simplest spiking neuron models widely used in SNNs. The membrane potential u(t) of the LIF neuron is defined as:

where Ii(t) denotes the input current from the i-th presynaptic neuron at time t, wi denotes the corresponding synaptic weight (connection weight). τm is the membrane time constant, and urest is the resting potential. When the membrane potential exceeds the firing threshold uth at time tf , a spike is generated and transmitted to postsynaptic neurons. At the same time, the membrane potential u(t) of the LIF neuron goes back to the reset value urest < uth (see Eq. (1)).

For numerical simulations of LIF neurons in SNNs, we need to consider a version of the parameters dynamics that is discrete in time. Specifically, Eq. (1) is approximated by:

Here mt and ut denote the membrane potential after neural dynamics and after the trigger of a spike at time-step t, respectively. St is the binary spike at time-step t, which equals

1 if there is a spike and 0 otherwise. H(·) is the Heaviside step function, which is defined by H(x) = {1 x ≥ 0 0 x < 0 . Note that H(0) = 1 is specially defined.

2.2 Deep Spiking Neural Network

A typical architecture of deep SNNs is exemplified as Fig. 1, which consists of neuronal layers (LIF layers) generating spikes and synaptic layers including convolutional layers, pooling layers, and fully-connected (FC) layers. The synaptic layers define connective patterns of synapses between spiking neurons, which are leveraged for spatial correlative feature extraction or classification. A deep SNN such as CIFARNet proposed in [Wu et al., 2019b] usually has multiple largescale synaptic layers and requires pruning for reducing computation cost.

2.3 Gradient Rewiring Framework

Inspired by synaptogenesis and synapse elimination of the brain, we present a gradient-based rewiring framework for deep SNNs pruning and call it Gradient Rewiring. Rewiring denotes the structural change of the brain's wiring diagram. Gradient Rewiring combines pruning and learning to explore both connective patterns and weight spaces. The combination of pruning and learning requires continuous parameterized expression of wiring change (synaptic connection change) and weight optimization. Previous neuroscience research has revealed that the wiring plasticity of neurons can be viewed as a special case of weight plasticity [Chklovskii et al., 2004]. Thus we can unify the representation of synaptic connectivity and synaptic weights by introducing a new hidden synaptic parameter θ in addition to the synaptic weight w.

The synaptic parameter θ is related to the volume of a dendritic spine, which defines the synaptic weights and connectivity for each synapse. A synapse is considered connected if θ > 0 and disconnected otherwise. Similar to [Bellec et al., 2018a], the synaptic parameter θ is mapped to synaptic weight through the relation w = sReLU(θ), where s ∈ {+1,-1} is fixed and used to distinguishes excitatory and inhibitory synapses. In this way, the sign change of synaptic parameter θ can describe both formation and elimination of synaptic connection.

Optimizing Synaptic Parameters θ

With the above definition, the gradient descent algorithm and its variants can be directly applied to optimize all synaptic parameters θ of the SNNs. Specifically, we assume that the loss function has the form L(xi, yi,w), where xi, yi denotes the i-th input sample and the corresponding target output, and w is the vector of all synaptic weights. The corresponding vectors of all synaptic parameters and signs are θ, s respectively. Suppose wk, θk, sk stand for the k-th element of w, θ, s. Then we have w = s ⊙ ReLU(θ), where ⊙ is the Hadamard product, and ReLU is an element-wise function. The gradient descent algorithm would adjust all the synaptic parameters to lower the loss function L, that is:

where H(·) represents the Heaviside step function. Ideally, if θk < 0 during learning, the corresponding synaptic connection will eliminate. It will regrow when θk > 0. Nonetheless, the pruned connection will never regrow. To see this, supposing that θk = θ < 0, we get H(θk) = 0, and wk = skReLU(θk) = 0. Then:

As the gradient of the loss function L with respect to θk is 0, θk will keep less than 0, and the corresponding synaptic connection will remain in the state of elimination. This might result in a disastrous loss of expression capabilities for SNNs. So far, the current theories and analyses can also systematically explain the results of deep rewiring [Bellec et al., 2018a], where the pruned connection is hard to regrow even noise is added to the gradient. There is no preference for minimizing loss when adding noise to gradient.

Gradient Rewiring

Our new algorithm adjusts θ in the same way as gradient descent, but for each θk, it uses a new formula to redefine the gradient when θk < 0.

For an active connection θk = θ > 0, the definition makes no difference as H(θk) = 1. For an inactive connection θk = θ < 0, the gradients are different as ∆θk = 0 and (∆θk)GradR ∝ sk ∂L ∂wk ∣∣wk=0 . The change of gradient for inactive connections results in a different loss function ¯ L. As illustrated in Fig. 2(a), when sk ∂L ∂wk ∣∣wk=0 < 0, the gradient of ¯ L guides a way to jump out the plateau in L (red arrow). With the appropriate step size chosen, a depleted connection will regrow without increasing the real loss L on this occasion. It should be noted that the situation in Fig. 2(b) can also occur and push the parameter θk even more negative, which can be overcome by adding a constraint scope to the parameter θ (discussed in the subsequent part).

Similar to the gradient descent algorithm (Eq. (3)), Grad R can also optimize synaptic parameters θ of the SNNs. An intuitive explanation is the angle between the accurate gradient ∇θL and the redefined gradient ∇θ ¯ L is below 90°, i.e., 〈∇θL,∇θ ¯ L〉 ≥ 0, which can be derived from Lemma 1. Further, a rigorous proof is illustrated in Theorem 1.

Theorem 1 (Convergence). For a spiking neural network with synaptic parameters θ, synaptic weights w, signs s, and the SoftPlus mapping, a smooth approximation to the ReLU, wk = sk β log(1+exp(βθk)), β ≫ 1, where wk, θk, sk stand for the k-th element of w, θ, s, the loss function L(θ) with lowerbound, if there exists a constant C such that ∀k, θk ≥ C and a Lipschitz constant L such that ‖∇θL(θ1)- ∇θL(θ2)‖ ≤ L‖θ1 - θ2‖ for any differentiable points θ1 and θ2, then L will decrease until convergence using gradient ∇θ ¯ L, if step size η > 0 is small enough.

Proof. Let θt be the parameter θ after t-th update steps using ∇θ ¯ L as direction for each step, that is, θt+1 = θt - η∇θ ¯ L(θt ). With Lipschitz continuous gradient, we have

where θt (τ) is defined by θt (τ) = θt + τ(θt+1 - θt ), τ ∈ [0, 1]. The first two inequality holds by Cauchy-Schwarz

inequality and Lipschitz continuous gradient, respectively. The last inequality holds with constant A > 0 according to Lemma 1.

When ‖∇θ ¯ L(θt )‖ > 0 and η < 2/AL, the last term in Eq. (6) is strictly negative. Hence, we have a strictly decreasing sequence {L(θn )}n∈N. Since the loss function has lowerbound, this sequence must converges.

Prior-based Connectivity Control

There is no explicit control of connectivity in the above algorithm. We need to explore speed control of pruning targeting diverse levels of sparsity. Accordingly, we introduce a prior distribution to impose constraints on the range of parameters. To reduce the number of extra hyperparameters, we adopt the same Laplacian distribution for each synaptic parameter θk:

where α and µ are scale parameter and location parameter respectively. The corresponding penalty term is:

which is similar to ℓ1 -penalty with sign(·) denoting the sign function.

Here we illustrate how the prior influences network connectivity. We assume that all the synaptic parameters θ follow the same prior distribution in Eq. (7). Since the connection with the negative synaptic parameter (θk < 0) is pruned, the probability that a connection is pruned is:

As all the synaptic parameters follow the same prior distribution, Eq. (9) also represents the network's sparsity. Conversely, given target sparsity p of the SNNs, one can get the appropriate choice of location parameter µ with Eq. (9), which is related to the scale parameter α

Note that when p < 0.5, we have µ > 0, which indicates that the peak of the density function p(θk) will locate at positive θk. As the sign sk is chosen randomly and set fixed, the corresponding prior distribution p(wk) will be bimodal (explained in Appendix C) . In view of that the distribution of weight (bias excluded) in convolutional and FC layers of a directly trained SNN is usually zero-meaned unimodal (See Fig. 3), we assume p ≥ 0.5 in the following discussion.

The Grad R algorithm is given in Algorithm 1. It is important to notice that the target sparsity p is not bound to be reached. Since the gradient of the loss function, i.e., sk ∂L(xi ,yi ,θ) ∂wk , also contributes to the posterior distribution (line 7). Consequently, the choice of these hyperparameters is

Algorithm 1: Gradient Rewiring with SGD

more a reference rather than a guarantee on the final sparsity. The prior encourages synaptic parameters to gather around µ, wherein the weights are more likely to be zero. Meanwhile, α controls the speed of pruning. The penalty term prevent weights from growing too large, alleviating the concern in Fig. 2(b).

2.4 Detailed Learning Rule in SNNs

Here we apply Grad R to the SNN in Fig. 1 and derive the detailed learning rule. Specifically, the problem is to compute the gradient of the loss function sk ∂L(xi ,yi ,θ) ∂wk (line 7 in Algorithm 1). We take a single FC layer without bias parameters as an example. Suppose that mt and ut represent the membrane potential of all the LIF neurons after neural dynamics and after the trigger of a spike at time-step t, St denotes the vector of the binary spike of all the LIF neurons at time-step t, W denotes the weight matrix of FC layer, 1 denotes a vector composed entirely of 1 with the same size as St. With these definitions, we can compute the gradient of output spikes St with respect to W by unfolding the network over the simulation time-steps, which is similar to the idea of backpropagation through time (BPTT). We obtain

When calculating ∂St ∂mt , the derivative of spike with respect to the membrane potential after charging, we adopt the surrogate gradient method [Wu et al., 2018; Neftci et al., 2019]. Specifically, we utilize shifted ArcTan function H1(x) = 1π arctan(πx) + 12 as a surrogate function of the Heaviside step function H(·), and we have ∂St ∂mt ≈H′1 (mt -uth).

3 Experiments

We evaluate the performance of the Grad R algorithm on image recognition benchmarks, namely the MNIST and CIFAR10 datasets. Our implementation of deep SNN are based on our open-source SNN framework SpikingJelly [Fang et al., 2020a].

3.1 Experimental Settings

For MNIST, we consider a shallow fully-connected network (Input-Flatten-800-LIF-10-LIF) similar to previous works. On the CIFAR-10 dataset, we use a convolutional SNN with 6 convolutional layers and 2 FC layers (Input-[[256C3S1-BN-LIF]×3-MaxPool2d2S2]×2-FlattenDropout-2048-LIF-100-LIF-AvgPool1d10S10, illustrated in Fig. 1) proposed in our previous work [Fang et al., 2020b]. The first convolutional layer acts as a learnable spike encoder to transform the input static images to spikes. Through this way, we elude from the Poisson encoder that requires a long simulation time-step to restore inputs accurately [Wu et al., 2019b]. The dropout mask is initiated at the first time-step and remains constant within an iteration of mini-batch data, guaranteeing a fixed connective pattern in each iteration [Lee et al., 2020]. All bias parameters are omitted except for BatchNorm (BN) layers. Note that pruning weights in BN layers is equivalent to pruning all the corresponding input channels, which will lead to an unnecessary decline in the expressiveness of SNNs. Moreover, the peak of the probability density function of weights in BN layers are far from zero (see Fig. 3), which suggest the pruning is improper. Based on the above reasons, we exclude all BN layers from pruning in our algorithm.

To facilitate the learning process, we use the Adam optimizer along with our algorithm for joint optimization of connectivity and weight. Furthermore, the gradients of reset ∂ut ∂St in LIF neurons' forward steps (Eq. 2) are dropped, i.e. manually set to zero, according to [Zenke and Vogels, 2020]. The choice of all hyperparameters is shown in Table 3 and Table 4.

3.2 Performance

Fig. 4 shows the top-1 accuracy on the test set and connectivity during training on the CIFAR-10 datasets. The improvement of accuracy and network sparsity is simultaneous as the Grad R algorithm can learn the connectivity and weights of the SNN jointly. As the training goes into later stages, the growth of sparsity has little impact on the performance. The sparsity will likely continue to increase, and the accuracy will still maintain, even if we increase the training epochs. The effect of prior is also clearly shown in Fig. 4. Overall, a larger penalty term α will bring higher sparsity and more apparent performance degradation. Aggressive pruning setting like huge α, e.g. α = 0.005, will soon prune most of the connections and let the network leaves few parameters to be updated (see the brown line).

As most of the previous works do not conduct experiments on datasets containing complicated patterns, such as CIFAR10, we also test the algorithm on a similar shallow fullyconnected network for the classification task of the MNIST dataset. We further implement Deep R as a baseline, which also can be viewed as a pruning-while-learning algorithm. The highest performance on the test set against connectivity is presented in Fig. 5. The performance of Grad R is comparable to Deep R on MNIST and much better than Deep R on CIFAR-10. As illustrated in Fig. 5(b), the performance of Grad R is close to the performance of the fully connected network when the connectivity is above 10%. Besides, it reaches a classification accuracy of 92.03% (merely <1% accuracy loss) when constrained to ∼5% connectivity. Further, Grad R achieves a classification accuracy of 89.32% (∼3.5% accuracy loss) under unprecedented 0.73% connectivity. These results suggest that there exists extremely high redundancy in deep SNNs. We compare our algorithm to other state-ofthe-art pruning methods of SNNs, and the results are listed in Table 1 and Table 2, from which we can find that the proposed method outperforms previous works.

Figure 5: Top-1 accuracies on test set. Different penalty term α is chosen for different sparsity (connectivity). The accuracy of unpruned model is shown as a black dashed line.

3.3 Behaviors of Algorithm

Ablation study of prior. We test the effects of prior. The Grad R algorithm without prior (α = 0) achieves almost the highest accuracy after pruning, but the connectivity drops the most slowly (the red curve in Fig. 4). The performance gap between unpruned SNN and pruned SNN without prior suggests the minimum accuracy loss using Grad R. It also clarifies that prior is needed if higher sparsity is required.

Distribution drift of parameters. Fig. 6(a) illustrates the distribution of the parameter θ during training (α = 0.005).

Since we use a Laplacian distribution with µ < 0 as prior, the posterior distribution will be dragged to fit the prior.

Competition between pruning and regrowth. A cursory glance of the increasing sparsity over time may mislead one that the pruning process always gains huge ascendancy over the regrowth process. However, Fig. 6(b) sheds light on the fierce competition between these two processes. The pruning takes full advantage of the regrowth merely in the early stages (the first 200 epochs). Afterward, the pruning process only wins by a narrow margin. The competition encourages SNNs to explore different network structures during training.

Firing rate and connectivity. Reduction of firing rate is essential for energy-efficient neuromorphic chips (event-driven hardware) [Oster et al., 2007]. Therefore, we explore the impact of synaptic pruning on the firing rate. As shown in Fig. 7, there is a relatively weak correlation between the distribution of neurons' average firing rate and connectivity. However, when the connectivity drastically drops extremely low under 1%, the level of neural activities is severely suppressed. The peak on the left side has a significant shift, causing by silent neurons. Meanwhile, the right peak representing saturated neurons vanishes.

Figure 7: The distribution of the average firing rate of post-synaptic neurons in the 1st FC layer under different connectivity. The firing rate of each neuron is averaged on samples of the whole test dataset and on time domain.

4 Conclusion

In this paper, we present Gradient Rewiring, a pruning method for deep SNNs, enabling efficient joint learning of connectivity and network weight. It utilizes the competition between pruning and regrowth to strengthen the exploration of network structures and achieves the minimal loss of SNN's performance on the MNIST and CIFAR-10 dataset so far. Our work reveals that there exists high redundancy in deep SNNs. An interesting future direction is to further investigate and exploit such redundancy to achieve low-power computation on neuromorphic hardware.

Acknowledgements

This work is supported by grants from the National Natural Science Foundation of China under contract No. 62027804, No. 61825101 and No. 62088102.

A Supporting Lemma

Lemma 1. If there exists a constant C, such that ∀k, θk ≥ C, then there exists a constant A > 0 such that

Proof. Since wk = sk β log(1 + exp(βθk)), we have ∂wk ∂θk = skσ(βθk), where σ(x) = 1 1+exp(-x) is sigmoid function.

Note that σ(x) > 0 is monotonically increasing, we have

Thus there exists a constant A = 1 σ(βC) such that ‖∇θ ¯ L‖2 ≤

Lemma 1 also suggests that 〈∇θL,∇θ ¯ L〉 ≥ 0. When ‖∇θ ¯ L(θt )‖ > 0, the inner product 〈∇θL,∇θ ¯ L〉 is strictly positive, thus the angle between the accurate gradient ∇θL and the redefined gradient ∇θ ¯ L is below 90°.

2 × 10-1

Table 4: Choice of penalty term for Grad R (target sparsity p = 0.95) and Deep R (maximum sparsity p = 1 for CIFAR-10, p = 0.99 for MNIST) with corresponding final connectivity.

B Detailed Setting of Hyperparameters

The hyperparameters of Grad R and our implementation of Deep R are illustrated in Table 3 and Table 4 for reproducing results in this paper. Note that for Deep R, zero penalty still introduces low connectivity.

C Range of Target Sparsity

Recall that the network weights w are determined by corresponding signs s and synaptic parameters θ, if the signs are randomly assigned and the distribution of θk is given, we are able to deduce the distribution of wk. For this algorithm, we denote the cumulative distribution function (CDF) of θk is Fθk (θ) := P[θk < θ]. So we obtain

Since the distribution of sign is similar to the Bernoulli distribution, the CDF of wk can also be derived as

Note that F ′wk (w) = 12 F ′θk (|w|), which indicates the probability density function (PDF) of wk is an even function. Therefore, the distribution of w is symmetric with respect to the w-axis. We hope it to be unimodal, which matches our observation of weight distribution in SNNs (See Fig. 3). So there should not be any peak in F ′wk (w) when w > 0, otherwise we will get the other peak symmetric to it.

Note that the requirement that F ′wk (w) = 12 F ′θk (|w|)(w > 0) has no peak is equivalent to that F ′θk (w)(w > 0) has no peak. As F ′θk (·) is the prior of θk, which is a Laplacian distribution where the PDF has only one peak at µ, we must constrain µ ≤ 0. This explains why we require target sparsity p ≥ 0.5.