云计算百科
云计算领域专业知识百科平台

时空注意力聚合器的脉冲神经网络(STAA-SNN)

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

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

在这里插入图片描述

一、研究背景

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

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

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

  • 三、本文创新点

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

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

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

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

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

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

  • 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 实现。

    六、结果

  • 静态图像
     • 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)。

  • 事件相机
     • CIFAR10-DVS:VGG-13 + T=16 达 82.10%(↑1.2% 超 Spikeformer)。
     • DVS128Gesture:VGG-13 + T=16 达 98.61%,与 LIAF+TA 持平但时间步减少 44%。

  • 消融实验(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× 的同时保持精度。

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

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

  • 七、意义与展望

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

    评论 抢沙发

    评论前必须登录!