Error-Aware Conversion from ANN to SNN via Post-training Parameter Calibration

Abstract

Spiking Neural Network (SNN), originating from the neural behavior in biology, has been recognized as one of the next-generation neural networks. Conventionally, SNNs can be obtained by converting from pre-trained Artificial Neural Networks (ANNs) by replacing the non-linear activation with spiking neurons without changing the parameters. In this work, we argue that simply copying and pasting the weights of ANN to SNN inevitably results in activation mismatch, especially for ANNs that are trained with batch normalization (BN) layers. To tackle the activation mismatch issue, we first provide a theoretical analysis by decomposing local layer-wise conversion error, and then quantitatively measure how this error propagates throughout the layers using the second-order analysis. Motivated by the theoretical results, we propose a set of layer-wise parameter calibration algorithms, which adjusts the parameters to minimize the activation mismatch. To further remove the dependency on data, we propose a privacy-preserving conversion regime by distilling synthetic data from source ANN and using it to calibrate the SNN. Extensive experiments for the proposed algorithms are performed on modern architectures and large-scale tasks including ImageNet classification and MS COCO detection. We demonstrate that our method can handle the SNN conversion and effectively preserve high accuracy even in 32 time steps. For example, our calibration algorithms can increase up to 63% accuracy when converting MobileNet against baselines.

Appendices

Appendix A Major Proof

In order to prove Theorem 4.1, we need to introduce two lemmas.

Lemma A.1

The error term $\textbf{e}_r^{({n})}$ can be computed using the error from former layer as

$$\begin{aligned} \textbf{e}_r^{(n)} = {\textbf{B}^{(n)}\textbf{W}^{(n-1)}\textbf{e}^{(n-1)}}, \end{aligned}$$

(A1)

where ${\textbf{B}^{(n)}}$ is a diagonal matrix whose values are the derivative of ReLU activation.

Proof

Without loss of generality, we will show that the equation holds for any two consecutive layers $(\ell +1)$ and $(\ell )$. Recall that the error term $\textbf{e}_c^{(\ell +1)}$ is the error between different input activation. Suppose $f(\textbf{a})=\textrm{ReLU}(\textbf{Wa})$ and $\textbf{e}_r^{(\ell +1)}=f(\textbf{x}^{(\ell )}) - f(\bar{\textbf{s}}^{(\ell )})$, we can rewrite $\textbf{e}_r^{(\ell +1)}$ also using the Taylor expansion, give by

$$\begin{aligned} \textbf{e}_r^{(\ell +1)} = \textbf{e}^{(\ell )}\nabla _{\textbf{x}^{(\ell )}}f + \frac{1}{2}\textbf{e}^{(\ell ),\top } \nabla ^2_{\textbf{x}^{(\ell )}}f\textbf{e}^{(\ell )} + {\textrm{O}(||\textbf{e}^{(\ell )}||^3)}, \end{aligned}$$

(2)

where $\textbf{e}^{(\ell )}$ is the difference between $\textbf{x}^{(\ell )}$ and $\bar{\textbf{s}}^{(\ell )}$. For the first order derivative, we easily write it as

$$\begin{aligned} \nabla _{\textbf{x}^{(\ell )}}f = {\textbf{B}^{(\ell +1)}\textbf{W}^{(\ell )}}, \end{aligned}$$

(3)

where $\textbf{B}^{(\ell +1)}\in \{0, 1\}^{c_{\ell +1}\times c_{\ell +1}}$ is a diagonal matrix, each element on the diagonal is the derivative of the ReLU function. $c_{\ell +1}$ denotes the number of neurons in layer $\ell +1$. To calculate the second order derivative, we need to differentiate the above equation with respect to $\textbf{x}^{(\ell )}$. First, the matrix $\textbf{B}^{(\ell +1)}$ is not a function of $\textbf{x}^{(\ell )}$ since it consists of only constants. Second, weight parameters $\textbf{W}^{(\ell )}$ also is not a function of $\textbf{x}^{(\ell )}$. Therefore, we can safely ignore the second-order term and other higher order terms in Eq. (2). Thus, Eq. (A1) holds. $\square$

Lemma A.2

((Botev et al., 2017) Eq. (8)) The Hessian matrix of activation in layer $(\ell )$ can be recursively computed by

$$\begin{aligned} \textbf{H}^{(\ell )} = {\textbf{W}^{(\ell ), \top }\textbf{B}^{(\ell +1), \top }\textbf{H}^{(\ell +1)}\textbf{B}^{(\ell +1)}\textbf{W}^{(\ell )}}. \end{aligned}$$

(4)

Proof

Let $\textbf{H}_{a, b}^{(\ell )}$ be the Hessian matrix (loss L with respect to $\ell$-th layer activation). We can calculate it by

$$\begin{aligned} \textbf{H}_{a, b}^{(\ell )}&= \frac{\partial ^2L}{\partial \textbf{x}_{b}^{(\ell )}\partial \textbf{x}_{a}^{(\ell )}}=\frac{\partial }{\partial \textbf{x}_{b}^{(\ell )}}\left( \sum _i \frac{\partial L}{\partial \textbf{x}_{i}^{(\ell +1)}}\frac{\partial \textbf{x}_{i}^{(\ell +1)}}{\partial \textbf{x}_{a}^{(\ell )}} \right) \end{aligned}$$

(5)

$$\begin{aligned}&= \sum _i \frac{\partial }{\partial \textbf{x}_{b}^{(\ell )}} \left( \frac{\partial L}{\partial \textbf{x}_{i}^{(\ell +1)}} \frac{\partial \textbf{x}_{i}^{(\ell +1)}}{\partial \textbf{z}_{i}^{(\ell )}} \frac{\partial \textbf{z}_{i}^{(\ell )}}{\partial \textbf{x}_{a}^{(\ell )}}\right) \end{aligned}$$

(6)

$$\begin{aligned}&= \sum _i \textbf{W}_{i,a}^{(\ell )} \frac{\partial }{\partial \textbf{x}_{b}^{(\ell )}} \left( \frac{\partial L}{\partial \textbf{x}_{i}^{(\ell +1)}} \frac{\partial \textbf{x}_{i}^{(\ell +1)}}{\partial \textbf{z}_{i}^{(\ell )}} \right) . \end{aligned}$$

(7)

Note that term $\frac{\partial \textbf{x}_{i}^{(\ell +1)}}{\partial \textbf{z}_{i}^{(\ell )}}$ is the derivative of the ReLU function, which is a constant and can be either 0 or 1, therefore it has no gradient with respect to $\textbf{x}_{b}^{(\ell )}$. Further differentiating the above equation, we have

$$\begin{aligned} \textbf{H}_{a, b}^{(\ell )}&= \sum _i \textbf{W}_{i,a}^{(\ell )} \left( \frac{\partial \textbf{x}_{i}^{(\ell +1)}}{\partial \textbf{z}_{i}^{(\ell )}} \frac{\partial ^2L}{\partial \textbf{x}_{b}^{(\ell )}\partial \textbf{x}_{i}^{(\ell +1)}}\right) \end{aligned}$$

(8)

$$\begin{aligned}&= \sum _i \textbf{W}_{i,a}^{(\ell )} \left[ \frac{\partial \textbf{x}_{i}^{(\ell +1)}}{\partial \textbf{z}_{i}^{(\ell )}} \left( \sum _j \frac{\partial ^2L}{\partial \textbf{x}_{j}^{(\ell +1)}\partial \textbf{x}_{i}^{(\ell +1)}}\frac{\partial \textbf{x}_{j}^{(\ell +1)}}{\partial \textbf{x}_{b}^{(\ell )}}\right) \right] \end{aligned}$$

(9)

$$\begin{aligned}&= \sum _i \textbf{W}_{i,a}^{(\ell )} \left[ \frac{\partial \textbf{x}_{i}^{(\ell +1)}}{\partial \textbf{z}_{i}^{(\ell )}} \left( \sum _j \textbf{W}_{j,b}^{(\ell )} \frac{\partial ^2L}{\partial \textbf{x}_{j}^{(\ell +1)}\partial \textbf{x}_{i}^{(\ell +1)}} \frac{\partial \textbf{x}_{j}^{(\ell +1)}}{\partial \textbf{z}_{j}^{(\ell )}} \right) \right] \end{aligned}$$

(10)

$$\begin{aligned}&= \sum _{i, j} \textbf{W}_{i,a}^{(\ell )} \frac{\partial \textbf{x}_{i}^{(\ell +1)}}{\partial \textbf{z}_{i}^{(\ell )}} \frac{\partial ^2L}{\partial \textbf{x}_{j}^{(\ell +1)}\partial \textbf{x}_{i}^{(\ell +1)}} \frac{\partial \textbf{x}_{j}^{(\ell +1)}}{\partial \textbf{z}_{j}^{(\ell )}} \textbf{W}_{j,b}^{(\ell )}. \end{aligned}$$

(11)

To this end, we can rewrite it using matrix form. Thus, the Lemma holds. $\square$

Now, we can prove our theorem with the above two lemmas.

Proof

According to Lemma A.1, the error in the last layer can be rewritten by

$$\begin{aligned} \textbf{e}^{(n)} = {\textbf{B}^{(n)}\textbf{W}^{(n-1)}\textbf{e}^{(n-1)}} + \textbf{e}_c^{(n)}. \end{aligned}$$

(12)

Applying this equation to the second-order objective, we have

$$\begin{aligned} \textbf{e}^{(n), \top }\textbf{H}^{(n)}\textbf{e}^{(n)}&= {\textbf{e}^{(n\!-\!1), \top }\textbf{W}^{(n\!-\!1), \top }\textbf{B}^{(n), \top } \textbf{H}^{(n)}\textbf{B}^{(n)}\textbf{W}^{(n\!-\!1)}\textbf{e}^{(n\!-\!1)} }\nonumber \\&\quad + {2\textbf{e}^{(n{-}1), \top }\textbf{W}^{(n{-}1), \top } \textbf{B}^{(n), \top }\textbf{H}^{(n)}\textbf{e}_c^{(n)}} \nonumber \\&\quad + \textbf{e}_c^{(n), \top }\textbf{H}^{(n)}\textbf{e}_c^{(n)}. \end{aligned}$$

(13)

According to Lemma A.2, the Hessian of activation can be derived recursively as

$$\begin{aligned} \textbf{H}^{(n-1)} = {\textbf{W}^{(n-1), \top }\textbf{B}^{(n), \top }\textbf{H}^{(n)}\textbf{B}^{(n)}\textbf{W}^{(n-1)}}. \end{aligned}$$

(14)

Thus, the first term in Eq. (13) can be rewritten to $\textbf{e}^{(n-1), \top }$$\textbf{H}^{(n-1)}\textbf{e}^{(n-1)}$. For the second term, we can upper bound it using Cauchy-Schwarz inequality $(x^\top Ay \le \sqrt{x^\top A xy^\top A y})$ and inequality of arithmetic and geometric mean $(\sqrt{x^\top A xy^\top A y}) \le \frac{1}{2}(x^\top A x+y^\top A y))$ by treating A as $\textbf{H}^{(n)}$. Therefore, the second term is upper bounded by

$$\begin{aligned} {2\textbf{e}^{(n-1), \top }}&{\textbf{W}^{(n-1), \top } \textbf{B}^{(n), \top }\textbf{H}^{(n)}\textbf{e}_c^{(n)}} \le \textbf{e}_c^{(n), \top }\textbf{H}^{(n)}\textbf{e}_c^{(n)} \nonumber \\ {}&+ {\textbf{e}^{(n-1), \top }\textbf{W}^{(n-1), \top }\textbf{B}^{(n), \top } \textbf{H}^{(n)}\textbf{B}^{(n)}\textbf{W}^{(n-1)}\textbf{e}^{(n-1)} }. \end{aligned}$$

(15)

To this end, we can rewrite Eq. (13) as

$$\begin{aligned} \textbf{e}^{(n), \top }\textbf{H}^{(n)}\textbf{e}^{(n)}&\le 2\textbf{e}^{(n-1), \top }\textbf{H}^{(n-1)}\textbf{e}^{(n-1)} + 2\textbf{e}_c^{(n), \top }\textbf{H}^{(n)}\textbf{e}_c^{(n)} \nonumber \\&\le \sum _{\ell =1}^n 2^{n-\ell +1} \textbf{e}_c^{(\ell ), \top }\textbf{H}^{(\ell )}\textbf{e}_c^{(\ell )}. \end{aligned}$$

(16)

$\square$

Appendix B Generalization to Under/Over-Fire Error

Here, we show that our analysis of error propagation and the calibration algorithm can generalize to under-fire and over-fire errors. Recall that in the case when $\textbf{v}^{\ell }(T)$ is not in $\left[ 0, V^{(\ell )}_{th}\right)$, under-fire and over-fire happen and we could not use a closed-form expression for $\bar{\textbf{s}}^{(\ell +1)}$. Here, we use function $\textrm{AvgIF}$ to denote the function of the averaged spike from the IF neuron:

$$\begin{aligned} \bar{\textbf{s}}^{(\ell +1)} = \textrm{AvgIF}(\textbf{W}^{\ell }\bar{\textbf{s}}^{(\ell )}). \end{aligned}$$

(B.17)

The $\mathrm {AvgIF(\textbf{W}^{\ell }\bar{\textbf{s}}^{(\ell )})}$ returns $\frac{m}{T}V_{th}^{\ell }$ where m is the number of output spikes. Note that with under/over-fire error the m cannot be analytically computed. Similarly, we define $\textbf{e}_c^{\ell }=\textrm{ReLU}(\textbf{W}^{\ell }\bar{\textbf{s}}^{(\ell )})-\textrm{AvgIF}(\textbf{W}^{\ell }\bar{\textbf{s}}^{(\ell )})$ as the local conversion error, which contains clipping error, flooring error, over-fire error, and under-fire error. Then, we can rewrite Eq. (13) as

$$\begin{aligned} \textbf{e}^{(n)} = \textbf{x}^{(n)}-\bar{\textbf{s}}^{(n)}&= \textrm{ReLU}(\textbf{W}^{(n-1)}\textbf{x}^{(n-1)}) \nonumber \\ {}&\quad - \textrm{AvgIF}(\textbf{W}^{(n-1)}\bar{\textbf{s}}^{(n-1)}) \nonumber \\&= \textrm{ReLU}(\textbf{W}^{(n-1)}\textbf{x}^{(n-1)}) \nonumber \\ {}&\quad - \textrm{ReLU}(\textbf{W}^{(n-1)}\bar{\textbf{s}}^{(n-1)}) \nonumber \\&\quad + \textrm{ReLU}(\textbf{W}^{(n-1)}\bar{\textbf{s}}^{(n-1)}) \nonumber \\ {}&\quad - \textrm{AvgIF}(\textbf{W}^{(n-1)}\bar{\textbf{s}}^{(n-1)}) \end{aligned}$$

(B.18)

$$\begin{aligned}&= \textbf{e}_r^{(n)} + \textbf{e}_c^{(n)}. \end{aligned}$$

(B.19)

Here, the accumulated error $\textbf{e}_r^{(n)}$ is the same with Eq. (13) and can be eliminated using second-order analysis. As for $\textbf{e}_c$, it can be addressed by our parameter calibration algorithm since our method is a data-driven approach. During calibration, the actual $\textbf{e}_c$ in each layer is computed using $\bar{\textbf{s}}^{(\ell +1)}-\textbf{x}^{(\ell )}$. Therefore, all errors can be reduced by calibration.

Appendix C Experiments Details

1.1 C.1 ImageNet Pre-training

The ImageNet dataset  (Deng et al., 2009) contains 120 M training images and 50 k validation images. For training pre-processing, we random crop and resize the training images to 224 $\times$ 224. We additionally apply CollorJitter with brightness = 0.2, contrast = 0.2, saturation = 0.2, and hue = 0.1. For test images, they are center-cropped to the same size. For all architectures we tested, the Max Pooling layers are replaced to Average Pooling layers. The ResNet-34 contains a deep-stem layer (i.e., three 3$\times$3 conv. layers to replace the original 7 $\times$ 7 first conv. layer) as described in  (He et al., 2019). We use Stochastic Gradients Descent with a momentum of 0.9 as the optimizer. The learning rate is set to 0.1 and followed by a cosine decay schedule  (Loshchilov & Hutter, 2016). Weight decay is set to $10^{-4}$, and the networks are optimized for 120 epochs. We also apply label smooth  (Szegedy et al., 2016)(factor = 0.1) and EMA update with 0.999 decay rate to optimize the model. For the MobileNet pre-trained model, we download it from pytorchcv².

1.2 C.2 CIFAR Pre-training

The CIFAR 10 and CIFAR100 dataset  (Krizhevsky et al., 2010) contains 50k training images and 10k validation images. We set padding to 4 and randomly cropped the training images to 32 $\times$ 32. Other data augmentations include (1)random horizontal flip, (2) Cutout  (DeVries & Taylor, 2017) and (3) AutoAugment  (Cubuk et al., 2019). For ResNet-20, we follow prior works  (Han et al., 2020; Han & Roy, 2020) who modify the official network structures proposed in  (He et al., 2016) to make a fair comparison. The modified ResNet-20 contains 3 stages with 4x more channels. For VGG-16 without BN layers, we add Dropout with a 0.25 drop rate to regularize the network. For the model with BN layers, we use Stochastic Gradients Descent with a momentum of 0.9 as the optimizer. The learning rate is set to 0.1 and followed by a cosine decay schedule  (Loshchilov & Hutter, 2016). Weight decay is set to $5\times 10^{-4}$ and the networks are optimized for 300 epochs. For networks without BN layers, we set weight decay to $10^{-4}$ and learning rate to 0.005.

1.3 C.3 COCO Pre-training

COCO dataset  (Lin et al., 2014) contains 121,408 images, with 883k objection annotations, and 80 image classes. The median image resolution is 640 $\times$ 480. The backbone ResNet50 is pre-trained on ImageNet dataset. We train RetinaNet and Faster R-CNN with 12 epochs, using SGD with the learning rate of 0.000625, decayed at 8, 11-th epoch by a factor of 10. The momentum is 0.9 and we use linear warmup learning in the first epoch. The batch size per GPU is set to 2 and in total we use 8 GPUs for pre-training. The weight decay is set to $1e-4$. To stabilize the training, we freeze the update in the BN layers (as they have low performance in small batch-size training), and the update in the first stage and stem layer of the backbone.

1.4 C.4 COCO Conversion Details

For the calibration dataset, we use 128 training images taken from the MS COCO dataset for calibration. The image resolution is set to 800 (max size 1333). Note that in object detection the image resolution is varying image-by-image. Therefore we cannot set an initial membrane potential at a space dimension (which requires a fixed dimension). Similar to bias calibration, we can compute the reduced mean value in each channel. Here we set the initial membrane potential as the $T\times \mu (\textbf{e}^{(i)})$. Learning hyper-parameters are kept the same with ImageNet experiments.