SEENN: 迈向时间脉冲早退神经网络的深度解析

Li Y, Geller T, Kim Y, et al. Seenn: Towards temporal spiking early exit neural networks[J]. Advances in Neural Information Processing Systems, 2023, 36: 63327-63342.

1. 引言与研究背景

脉冲神经网络（Spiking Neural Networks, SNNs）作为第三代神经网络，通过模拟生物神经元的脉冲发放机制，在能效和生物合理性方面展现出独特优势。与传统人工神经网络（ANNs）在单个时间步内处理信息不同，SNNs通过多个时间步的脉冲序列进行信息编码和处理。这种时空处理方式使得SNNs特别适合处理时序信息和事件驱动的数据，同时其二进制脉冲特性将乘法运算替换为加法运算，大幅降低了计算复杂度。

然而，SNNs面临一个根本性挑战：时间步数$T$与模型性能之间的权衡。增加时间步数能够捕获更多时间维度的特征，提高模型准确率，但同时增加了推理延迟和计算需求。现有方法通常将$T$设置为固定超参数，对所有输入样本使用相同的时间步数，忽略了样本间的难度差异。

2. 问题分析与动机

2.1 实验观察

论文通过大量实验揭示了一个关键现象：SNNs的准确率随时间步数增加而提升，但边际收益递减明显。

图1详细描述：该图展示了在多个数据集上SNNs的准确率-时间步曲线。图中包含四个子图：

• CIFAR-10（训练）：准确率从$T=2$时的94.16%增至$T=6$时的94.50%，仅提升0.34%，AET=2.228
• CIFAR-100（训练）：准确率从$T=2$时的72.87%增至$T=6$时的74.72%，提升1.85%，AET=3.053
• ImageNet（训练）：准确率从$T=1$时的60.75%增至$T=4$时的68.0%，AET=2.065
• CIFAR-10（转换）：展示了ANN-SNN转换方法的性能，AET=1.484

这些曲线清晰地表明，大多数样本可以在较少的时间步内被正确分类，而增加时间步主要是为了处理少数困难样本。

2.2 理论基础

基于上述观察，论文提出了一个核心假设：

假设3.1：给定一个脉冲神经网络$fT$，如果它能在$t$个时间步内正确预测输入$x$，那么对于任何$t' \in [t, T]$，$f{t'}(x)$也会输出正确预测。

这一假设表明不同时间步之间存在包含关系：$C_1 \subseteq C_2 \subseteq \cdots \subseteq C_T$，其中$C_t$表示在时间步$t$能够正确分类的样本集合。

3. 方法详述

3.1 脉冲神经网络的数学框架

SNNs的核心是泄漏积分发放（LIF）神经元模型。对于第$\ell$层的神经元，其动态演化过程由以下方程描述：

膜电位更新：
$$u^{\ell}[t+1] = \tau u^{\ell}[t] + W^{\ell}s^{\ell-1}[t]$$

脉冲发放：
$$s^{\ell}[t+1] = H(u^{\ell}[t+1] - V)$$

膜电位重置（如果发放脉冲）：
$$u^{\ell}[t+1] = u^{\ell}[t+1] \cdot (1 - s^{\ell}[t+1])$$

其中，$\tau \in (0,1]$是泄漏因子，控制膜电位的时间衰减；$W^{\ell}$是第$\ell$层的权重矩阵；$V$是发放阈值；$H(\cdot)$是Heaviside阶跃函数。

整个SNN的输出通过时间平均得到：
$$f_T(x) = \frac{1}{T}\sum_{t=1}^{T} h \circ g^L \circ g^{L-1} \circ \cdots \circ g^1(x)$$

3.2 平均最早时间步（AET）度量

为量化早退出的潜力，论文提出AET度量：

$$\text{AET} = \frac{1}{N}\left(|C_1| + T|W| + \sum_{t=2}^{T}t(|C_t| - |C_{t-1}|)\right)$$

该公式的推导基于以下考虑：

• $|C_1|$个样本在第1个时间步就能正确分类，贡献$|C_1| \times 1$
• 对于$t \geq 2$，新增的$(|Ct| - |C{t-1}|)$个样本需要$t$个时间步
• 错误分类的$|W|$个样本使用最大时间步$T$

3.3 SEENN-I：基于置信度的早退出

图2(a)详细描述：左侧展示了SEENN-I的架构。输入图像首先经过脉冲神经网络处理，在每个时间步产生输出。这些输出被累积并计算置信度分数（CS）。如果CS超过预设阈值$\alpha$，则使用当前累积输出进行预测；否则，继续下一个时间步的迭代。

置信度分数定义为预测概率分布的最大值：
$$\text{CS} = \max(p), \quad \text{其中} \quad p = \text{softmax}(f_t(x))$$

SEENN-I的推理算法如下：

1. 初始化累积输出$o = 0$
2. 对于$t = 1, 2, \ldots, T$：
   • 计算第$t$步输出：$o_t = h \circ g^L \circ \cdots \circ g^1(x)$
   • 更新累积输出：$o = o + o_t$
   • 计算置信度：$\text{CS} = \max(\text{softmax}(o/t))$
   • 如果$\text{CS} > \alpha$，返回$\arg\max(o/t)$并退出

3.4 SEENN-II：基于强化学习的早退出

图2(b)详细描述：右侧展示了SEENN-II的架构。策略网络接收输入图像，输出一个概率分布，用于采样时间步数。选定的时间步数用于SNN推理，产生预测结果。根据预测的正确性计算奖励，并通过策略梯度更新策略网络。

策略网络产生分类分布：
$$v = \text{softmax}(f_p(x;\theta))$$

其中$v = [v_1, v_2, \ldots, v_n]$，$v_k$表示选择$t_k$个时间步的概率。

动作采样遵循分类分布：
$$\pi_\theta(z|x) = \prod_{k=1}^{n}v_k^{z_k}$$

其中$z \in {0,1}^n$是one-hot向量，表示选择的时间步。

奖励函数设计平衡准确性和效率：
$$R(z) = \begin{cases} \frac{1}{2t_k|_{z_k=1}} & \text{预测正确} \\ -\beta & \text{预测错误} \end{cases}$$

3.5 训练策略

对于SEENN-I，采用时间有效训练（TET）损失：
$$\min_W \frac{1}{n}\sum_{k=1}^{n}\mathcal{L}_{CE}(f_{t_k}(x), W, y)$$

这确保每个时间步的预测都得到优化，提高早期时间步的预测质量。

对于SEENN-II，联合优化目标为：
$$\min_{W,\theta} \mathbb{E}_{z\sim\pi_\theta}\left[-R(z) + \mathcal{L}_{CE}(f_{t_k|_{z_k=1}}(x), W, y)\right]$$

4. 实验结果与分析

4.1 性能对比

表1 - 直接训练方法对比（CIFAR数据集）：

• ResNet-19在CIFAR-10上：SEENN-I用1.09个平均时间步达到96.07%准确率，而TEBN用2个时间步仅达95.45%
• VGG-16在CIFAR-100上：SEENN-II用1.15个平均时间步达到72.76%准确率，显著优于其他方法

表2 - ANN-SNN转换方法对比：

• ResNet-18在CIFAR-10上：SEENN-I用1.40个平均时间步达到93.63%准确率，而QCFS用4个时间步才达90.43%

4.2 硬件效率评估

图3详细描述：该图包含两个子图，分别展示吞吐量和能耗对比。左图显示SEENN-I和SEENN-II的吞吐量显著高于固定时间步SNN，特别是与4时间步SNN相比，吞吐量提升约3倍。右图展示能耗对比，SEENN方法在保持高准确率的同时，能耗降低约70%。

4.3 消融研究

图4详细描述：左侧展示了不同置信度阈值$\alpha$下的准确率-时间步曲线。随着$\alpha$从0.4增加到0.9，平均时间步从1.33增至2.35（ImageNet）和从1.06增至1.63（CIFAR-10）。右侧的饼图清晰展示了输入样本的时间步分布，例如在$\alpha=0.4$时，73.4%的样本仅需1个时间步。

4.4 定性分析

图5详细描述：该图展示了6个类别的代表性图像，每个类别包含简单（T=1）和困难（T=4）样本。简单样本通常具有清晰背景、单一主体、良好光照；困难样本则存在遮挡、多个对象、复杂背景或不利视角。例如，橙子类别中，T=1的图像是单个清晰的橙子，而T=4的图像中橙子与其他水果混在一起。

5. 附录：数学推导

A. 策略梯度

目标函数：
$$J(\theta) = \mathbb{E}_{z\sim\pi_\theta}[R(z)]$$

梯度计算：
$$\nabla_\theta J(\theta) = \mathbb{E}_{z\sim\pi_\theta}[R(z)\nabla_\theta\log\pi_\theta(z|x)]$$

由于$\pi\theta(z|x) = \prod{k=1}^n v_k^{z_k}$，我们有：
$$\log\pi_\theta(z|x) = \sum_{k=1}^n z_k\log v_k$$

因此：
$$\nabla_\theta\log\pi_\theta(z|x) = \sum_{k=1}^n z_k\nabla_\theta\log v_k$$

对于分类分布，期望可以精确计算：
$$\nabla_\theta J(\theta) = \sum_{k=1}^n R(z|_{z_k=1})v_k\nabla_\theta\log v_k$$

其中$R(z|_{z_k=1})$是使用$t_k$个时间步时的奖励。

B. 置信度阈值的理论分析

设预测概率分布为$p = [p_1, p_2, \ldots, p_M]$，其中$M$是类别数。置信度分数$\text{CS} = \max_i p_i$。

对于正确分类的样本，设真实类别为$y$，则：
$$P(\text{正确}|\text{CS} = c) = P(p_y = c | p_y = \max_i p_i)$$

根据贝叶斯定理和经验观察，这个概率随$c$单调递增，为阈值选择提供了理论依据。

C. 时间复杂度

设网络有$L$层，每层计算复杂度为$O(N^2)$（$N$是神经元数）。

固定时间步SNN的复杂度：
$$O(T \cdot L \cdot N^2)$$

SEENN的期望复杂度：
$$O(\bar{T} \cdot L \cdot N^2 + N_{policy})$$

其中$\bar{T}$是平均时间步数，$N{policy}$是策略网络的复杂度。由于$\bar{T} \ll T$且$N{policy} \ll L \cdot N^2$，SEENN显著降低了计算复杂度。

D. 收敛性

对于SEENN-II的策略梯度优化，在满足以下条件时保证收敛：

1. 学习率满足Robbins-Monro条件：$\sum_t \alpha_t = \infty$，$\sum_t \alpha_t^2 < \infty$
2. 奖励函数有界：$|R(z)| \leq R_{max}$
3. 策略网络Lipschitz连续

在这些条件下，策略梯度算法收敛到局部最优。

6. 结论

SEENN通过引入样本自适应的时间步机制，成功解决了SNNs中准确性与效率的固有权衡。实验结果表明，SEENN能够在保持甚至提升准确率的同时，将平均时间步数降低至原来的1/5到1/3，相应地减少了能耗和延迟。这种细粒度的时间维度控制为神经形态计算的实际部署开辟了新的可能性，特别是在边缘计算和实时处理场景中具有重要应用价值。