【学习笔记】Batch Norm in SNN

Batch Norm in SNN

对脉冲卷积神经网络进行基于时空域的标准化

Going Deeper With Directly-Trained Larger Spiking Neural Networks

背景

目前而言，行之有效的深度SNN训练方法大致有以下两种思路：

1. 将预先训练好的ANN转换为SNN版本
   –> 该方法一般需要较长的脉冲编码窗口来达到收敛的效果
2. 直接在时空域训练SNN
   –> 由于二元脉冲的激活函数问题以及梯度消失或爆炸问题导致网络层数不能过深

为解决后者存在的问题，提出了基于时空反向传播（spatio-temporal backpropagation, STBP)的，依赖于阈值的批标准化（threshold-dependent batch normalization method， tdBN）方法，称为"STBP-tdBN"，能够直接训练网络非常深的SNN，并在相应硬件上有效实现推断。

具体实现

对于一个脉冲卷积神经网络（Spiking Convolution Neural Network, SCNN），时间步长t内前向传播过程如下：
$$x_t=W⊛o^t+B$$
其中$x_t$为突触前输入信号（突触前膜输入信号的加权和，尚未经过IF模型产生输出脉冲），$x^t\in R^{N\times C\times H\times W}$。

在tdBN算法中，将突触前输入依据其通道作相应的归一化。假定$x_k^t$代表$x^t$特征图的第k维，则$x_k=(x_k^1,x_k^2,\cdots ,x_k^T)$将被标准化为：
$${\hat{x}}_k=\frac{\alpha V_{th}(x_k-E[x_k])}{\sqrt{Var[x_k]+ϵ}}$$

$$y_k={\lambda }_k{\hat{x}}_k+{\beta }_k$$

依定义有：
$$E[x_k]=mean(x_k)$$

$$Var[x_k]=mean((x_k-E[x_k]{)}^2)$$

P.S. 把每个timestep得到的featuremap中第k通道分别取出，合并得到$x_k$，后进行均值和方差的计算。

为避免“先卷积后BN”的操作增加不必要的计算量，可以利用batchnorm-scale-fusion的操作进行融合，从而将BN归并到卷积中。设若$W_{c,k}$和$B_{c,k}$，分别代表当前特征图c通道和下一特征图k通道之间的卷积核以及偏置项，则有：
$$W_{c,k}^{\prime }={\lambda }_k\frac{\alpha V_{th}{W}_{c,k}}{\sqrt{{\sigma }_{inf,k}^2+ϵ}}$$

$$B_{c,k}^{\prime }={\lambda }_k\frac{\alpha V_{th}(B_{c,k}-{\mu }_{inf,k})}{\sqrt{{\sigma }_{inf,k}^2+ϵ}}+{\beta }_k$$

简言之，tdBN与标准BN的主要区别有两处：

1. 有别于ANN，SNN不仅逐层传播，同时逐时刻传播信息。因此stBN应当在时间和空间两个维度上对特征输入归一化。
2. 归一化方差依赖于阈值的选择。在stBN中，pre-activations被标准化为$N(0,(\alpha V_{th}{)}^2$，（一般常用标准化为0-1正态分布）