STAA-SNN: Spatial-Temporal Attention Aggregator for Spiking Neural Networks

代码：https://github.com/GuoLab-UESTC/SCTFA-SNN

一、研究背景

1. 能耗与生物可解释性：SNN 用 0/1 脉冲事件驱动，乘-加运算大量替换为纯加运算，在 45 nm 工艺下 0.9 pJ/AC 比 4.6 pJ/MAC 节能 5× 以上，且与生物神经元“发放-整合-泄漏”行为一致。
2. 性能瓶颈：直接训练的大深度 SNN 在 CIFAR-100、ImageNet 上仍比同规模 ANN 低 3–10%，成为落地最大障碍。
3. 注意力机制在 ANN 中的成功，启发将其引入 SNN；但 SNN 的“时空混合、二值稀疏、多时间步迭代”三大特性，使传统注意力无法直接迁移。

二、现有方法局限

1. 训练方式
    • ANN→SNN 转换：需 64–256 时间步补偿量化误差，延迟高，且无法利用脉冲动态。
    • 无监督（STDP/Hebb）：难以做深层端到端优化。
    • 有监督 BPTT+替代梯度：梯度消失、膜电位震荡、易陷局部最优。

2. 注意力尝试
    • TA-SNN、MA-SNN、TCJA-SNN 等仅把“通道-空间-时间”拆成独立模块，未在单个神经元层内同时完成时空耦合；
    • 无位置编码，时间顺序信息靠 RNN 式递归隐式携带，深度增加后衰减严重；
    • 未解决“多步特征固化”——越深时间步越易早熟，导致网络停止更新。

三、本文创新点

1. 首次在 SNN 内部提出“脉冲驱动的自注意力”——用 Global Context (GC) 块实现纯卷积形式的 K/Q/V 自注意，避免 Transformer 的二次复杂度。
2. 首次把“位置编码（PE）”引入 SNN：显式把时间步索引嵌入到特征通道，弥补脉冲稀疏带来的时序信息丢失。
3. 提出“步注意力 SA”——对每个时间步的 4D（T,C,H,W）特征图重新加权，抑制噪声步、放大关键步，实现“时间维度软选择”。
4. 提出训练策略 TSRD：在 (0,T) 区间以概率 β 随机丢弃上述增强模块，回退到简单相加，防止早熟局部最优，加速收敛。
5. 整体即插即用：GC/PE/SA 三模块+TSRD 可嵌入任何基于 LIF 的 CNN 结构，不改变下游任务损失。

四、方法论细节

1. 自适应 LIF
    将原来固定泄漏系数 1/τ、重置电位 Vreset 改成可学习矩阵 M、N，让每层兴奋/抑制强度可自动调节。

2. GC 块（空间自注意）
    输入 Xt,n∈RC×H×W → 1×1 卷积生成 K/Q/V → 通过 Sigmoid+LN+ReLU 得到权重 Wv∈RC×1×1 → 与原特征残差相加，实现“脉冲-驱动”的自注意。

3. PE 块
    生成可学习向量 Post,n∈RC，在时间轴上与输入 It,n 相加，使网络知道“当前是第几步”。

4. SA 块（步注意）
    对每步特征 Ut,n 先 AvgPool 压缩→Conv1→ReLU→Conv2→Sigmoid 得到步权重，与 Ut,n 逐像素相乘，实现“哪一步更重要”的自适应加权。

5. STAA-LIF 前向公式
    Ut,n = GC1(Xt,n) ⊕ GC2(Ht-1,n) // 空间聚合
    Vt,n = SA(Ut,n)　　　　　　　　　// 时间加权
    再代入标准 LIF 更新膜电位与脉冲。

6. TSRD
    训练阶段以 β=0.1 概率把 GC/PE/SA 全部旁路，只做简单相加；推理阶段全部打开。类似 Dropout，但作用于“模块级”而非“神经元级”。

五、实验设置

数据集：CIFAR-10/100、ImageNet（静态）；CIFAR10-DVS、DVS128Gesture（事件相机）。
架构：ResNet-18/19/20/34、VGG-13，时间步 T=1/2/4/16。
基线：GLIF、TET、LSG、PFA、MPBN、IM-Loss、Spikformer 等 20 余条 SOTA。
硬件：单卡 RTX 3090（ImageNet 用 8×4090），PyTorch 实现。

六、结果

1. 静态图像
    • CIFAR-10：ResNet-19 + T=4 达 97.14%（↑0.42% 超 MPBN），且 T=1 即可 96.75%，比同量级方法少用 1–3 个时间步。
    • CIFAR-100：ResNet-19 + T=4 达 82.05%（↑1.95% 超 MPBN）。
    • ImageNet：ResNet-34 + T=4 达 70.40%（↑2.71% 超 SEW-ResNet）。

2. 事件相机
    • CIFAR10-DVS：VGG-13 + T=16 达 82.10%（↑1.2% 超 Spikeformer）。
    • DVS128Gesture：VGG-13 + T=16 达 98.61%，与 LIAF+TA 持平但时间步减少 44%。

3. 消融实验（CIFAR-100, T=4）
    Baseline
    → +GC +0.92%
    → +PE +0.57%
    → +SA +0.82%
    → +TSRD +0.32%，累计提升 2.80%。
    最佳 dropout 概率 β=0.1；GC 压缩系数 r=4 可在参数量↓4× 的同时保持精度。

4. 可视化
    Grad-CAM 显示 STAA-LIF 在 CIFAR10-DVS 马样本上聚焦马头、马嘴等判别区域，而 Vanilla LIF 仅粗略轮廓。

5. 能耗
    ResNet-20 在 CIFAR-100 上单张 32×32、T=4 仅 0.366 mJ，约为同精度 ANN 的 1/5。

七、意义与展望

1. 学术价值：首次把“空间-自注意 + 时序位置编码 + 步注意力”统一在脉冲神经元内，证明 SNN 也能具备 Transformer 类似的动态选择能力，而无需高功耗的密集矩阵乘法。
2. 工程价值：即插即用模块 + TSRD 训练策略，可直接嵌入现有 CNN-SNN 框架，在相同或更少时间步下获得显著精度提升，缩短推理延迟，降低边缘端能耗。
3. 未来方向
    • 将 GC 块进一步稀疏化，用事件驱动卷积核减少计算量；
    • 把 TSRD 与动态时间步自适应结合，实现“训练-推理时间步不一致”的极致压缩；
    • 在更具挑战的检测、分割、多模态任务上验证泛化能力；
    • 与 neuromorphic 芯片（Loihi、TrueNorth）联合部署，实测毫瓦级功耗下的精度-延迟权衡。