跳到主要内容
起居室老虎
返回

SpikeVideoFormer: An Efficient Spike-Driven Video Transformer with Hamming Attention and O(T) Complexity

2,030 字约 8 分钟

原文:arXiv:2505.10352 · Zou et al., ICML 2025

摘要: 脉冲神经网络在多种视觉任务中已表现出与人工神经网络相当的性能,同时具备更优的能效。然而,现有基于 SNN 的 Transformer 主要聚焦于单幅图像任务,强调空间特征,而未能充分利用 SNN 在视频视觉任务中的效率优势。本文提出了一种高效的脉冲驱动视频 Transformer——SpikeVideoFormer,其时间复杂度为线性 𝒪(T)𝒪(T)。具体而言,我们设计了脉冲驱动汉明注意力(Spike-Driven Hamming Attention, SDHA),在理论上实现了从传统实值注意力到脉冲驱动注意力的无缝转换。基于 SDHA,我们进一步分析了多种脉冲驱动的时空注意力结构,并确定了一种在保持线性时间复杂度的同时,能为视频任务提供优异性能的最优方案。SpikeVideoFormer 在视频分类、人体姿态跟踪和语义分割等多种下游视频任务上展现了出色的泛化能力与高效率。实验证明,与现有 SNN 方法相比,我们在后两项任务上取得了超过 15% 的性能提升;同时在保持与最新 ANN 方法相当精度的情况下,效率分别提升 ×16、×10 和 ×5。

1. Intro

主要讲 SNN 的 Energy Efficiency 来自于加法 & 稀疏。注意到现在的 SNN 主要是做单帧的静态输入,但是一直在 claim SNN 具有时空数据敏感的特征,因此认为 SNN 应当适合视觉任务。

SNN 迁移到 Video Task 上的主要 Challenge:

  1. 现在的 Spike Attention 集中沿用 ANN 中的点乘操作,不一定能够刻画 Spike 之间的相似性(见下图,比如两个 Spike Seq 之间因为 Spike 错开而形成 0,但两者不一定就完全不相关);
  2. 这些 Attention 都是 Space 上的而不是 Temporal 上的。

image.png

为了解决这个问题,文章的核心是提出了一种基于汉明距离的注意力机制,然后考虑汉明相似度存在浮点计算、不可微分、非线性复杂度问题,又对这种注意力机制进行了改造。

3. Method

3.1. Spiking Neuron Model

用的 LIF,记作SN(U)\text{SN}(U)

3.2. Spike-Driven Video Transformer

image.png

原始采样T×H×W×3T\times H\times W\times 3,首先经过下采样 + Spike 化,变成T×H8×W8×4CT\times \frac{H}{8}\times\frac{W}{8}\times 4C,其中的每一个 Conv Module 实际上都包括一个 DS Conv + 一个通道卷积:

U′=U+SepConv(U)SepConv(U)=Convpw(Convdw(SN(Convpw(SN(U)))))U′′=U′+ChannelConv(U′)ChannelConv(U′)=Conv(SN(Conv(SN(U′))))\begin{align*} U'&=U+\text{SepConv}(U)\\ \text{SepConv}(U)&=\text{Conv}_{\text{pw}}(\text{Conv}_\text{dw}(\text{SN}(\text{Conv}_\text{pw}(\text{SN}(U)))))\\ U''&=U'+\text{ChannelConv}(U')\\ \text{ChannelConv}(U')&=\text{Conv}(\text{SN}(\text{Conv}(\text{SN}(U')))) \end{align*}

Convpw\text{Conv}_\text{pw}:Point-wise 卷积;Convdw\text{Conv}_\text{dw}:Depth-wise Seperatable 卷积。

得到一个T×H16×W16×10CT\times \frac{H}{16}\times\frac{W}{16}\times 10C的输出,然后送入 Spike-Driven Spatiotemporal Transformer:

QS,KS,VS=Reshape(SN(RepConv(U)))U′=U+RepConv(SDHA(QS,KS,VS))U′′=U′+ChannelMLP(U′)ChannelMLP(U′)=MLP(SN(MLP(SN(U′))))\begin{align*} Q_S,K_S,V_S&=\text{Reshape}(\text{SN}(\text{RepConv}(U)))\\ U'&=U+\text{RepConv}(\text{SDHA}(Q_S, K_S, V_S))\\ U''&=U'+\text{ChannelMLP}(U')\\ \text{ChannelMLP}(U')&=\text{MLP}(\text{SN}(\text{MLP}(\text{SN}(U')))) \end{align*}

RepConv\text{RepConv}是重参数化卷积。

3.3. Spike-Driven Hamming Attention (SDHA)

image.png

脉冲查询向量与脉冲键向量在使用 点积 与 归一化汉明相似度 时的注意力分数直观对比。查询向量某位为 0 时,点积会忽略对应键向量的该位,导致四个差异极大的键向量得到相同得分,暴露出点积对二值脉冲向量相似度刻画的不适用性。

Spike-Driven Self-Attention (SDSA)

SDSA=𝒮𝒩s((QsKs⊤)Vs⏟𝒪(N2D))=𝒮𝒩s(Qs(Ks⊤Vs)⏟𝒪(ND2)),\mathrm{SDSA} = \mathcal{SN}_s\!\Bigl( \underbrace{(Q_s K_s^{\top})\,V_s}_{\mathcal{O}(N^{2}D)} \Bigr) \;=\; \mathcal{SN}_s\!\Bigl( Q_s\, \underbrace{(K_s^{\top} V_s)}_{\mathcal{O}(N D^{2})} \Bigr),

𝒮𝒩s\mathcal{SN}_s是带有 scaled threshold s⋅uths\cdot u_\text{th}的 LIF Neuron。

然而,点积并不适合作为脉冲注意力中的相似度度量。具体来说,当脉冲查询向量中含有零元素时,点积会无差别地忽略脉冲键向量对应位置的元素,容易造成特征消失或混淆。为直观展示这一问题,我们在图 2 中给出了两个示例:点积为四个差异显著的脉冲键向量得出了完全相同的分数。这说明 SDSA 所采用的点积无法准确反映两个二值脉冲向量之间的相似性。此外,式 (3) 中的阈值缩放因子ss只能凭经验设定,缺乏保证训练稳定性的明确指导。

Proposition 3.1: JL Lemma on Binary Embedding

设有:

𝐪,𝐤∈RC\textbf{q}, \textbf{k}\in R^C

分别是实数域上的 Query 和 Key 向量,对应的 Binary Embedding 为:

𝐪s,𝐤s∈{0,1}D\textbf{q}_s, \textbf{k}_s\in\{0, 1\}^D

定义为:

𝐪s=sign(A𝐪),𝐤s=sign(A𝐤)\textbf{q}_s=\text{sign}(A\textbf{q}), \textbf{k}_s=\text{sign}(A\textbf{k})

投影矩阵:

A∈RC×D,Aij∼N(0,1) i.i.dA\in R^{C\times D}, \quad A_{ij}\sim \text{N}(0, 1) \text{ i.i.d}

给定任意δ>0\delta>0且

D>log⁡Mδ2D>\frac{\log M}{\delta ^2}

有:

Pr(|fH(𝐪s,𝐤s)−g(fC(𝐪,𝐤)|≤δ)≥1−2e−δ2D\text{Pr}\left(|f_H(\textbf{q}_s, \textbf{k}_s)-g(f_C(\textbf{q}, \textbf k)|\le \delta \right)\ge 1-2e^{-\delta^2D}

上面的定理实际上证明了,在 channel size DD足够大的时候,汉明相似度以高概率逼近g(fc)g(f_c),而fcf_c是传统的余弦相似度。由于g(⋅)g(\cdot)是连续且单调的,它相当于保持了原有的 key-query 之间的相似度排序。

进一步的实验证明:

image.png

Computational Efficiency Design

直接使用fH(⋅)f_H(\cdot)的做法的问题:

  1. 1D\frac{1}{D}引入了浮点运算
  2. 有指示函数,不可微分无法反向传播
  3. 随着 token 增长,计算的复杂度不是线性增加的

为了解决上面的问题,将fHf_H重写为:

fh(𝐪s,𝐤s)=12+12D(2𝐪s−1)⊤(2𝐤s−1)f_h(\textbf q_s, \textbf k_s)=\frac{1}{2}+\frac{1}{2D}(2\textbf q_s-1)^\top(2\textbf k_s-1)

则:

SDHA=𝒮𝒩([12D(2Qs−𝟏)(2Ks−𝟏)⊤]Vs)\mathrm{SDHA} = \mathcal{SN}\!\Bigl( \bigl[\tfrac{1}{2D}\,(2Q_s-\mathbf 1)(2K_s-\mathbf 1)^{\!\top}\bigr]\,V_s \Bigr)

或者

SDHA=𝒮𝒩2D((2Qs−𝟏)[(2Ks−𝟏)⊤Vs]⏟𝒪(ND2))\mathrm{SDHA} = \mathcal{SN}_{2D}\!\Bigl( (2Q_s-\mathbf 1)\, \underbrace{\bigl[(2K_s-\mathbf 1)^{\!\top} V_s\bigr]}_{\mathcal{O}(N D^{2})} \Bigr)

可以把12D\frac{1}{2D}吸收到神经元的阈值中,用移位处理2Qs−12Q_s-1的计算,这样推理过程中不存在乘除法的计算。

3.4. Space-Time Spike-Driven Attention Designs

image.png

image.png

ANN 领域典型的做法如上。SpikeVideoFormer 采用联合注意力机制,将 feature 变成B×(TN)×DB\times (TN)\times D然后整个打进 Transformer。

4. Experiments

在 Video Classification, Pose Estimation , Video Semantic Segmentation 上做。

image.png

image.png

推理时间

在 A6000 GPU 与 AMD EPYC 7543 CPU 上测试批量为 1 和 1000 的视频平均耗时:

  • GLoT(ANN,二次 attention)当 T:8→32 时,时延 303→2972 ms(9.8×)
  • VIBE(ANN-GRU)时延 264→1335 ms(5.1×)
  • 本方法(SNN,线性 attention)时延 235→1087 ms(4.6×)

边缘部署

SNN 仅使用加法运算,更适合神经计算芯片部署,可显著降低功耗和时延。在 AMD Xilinx ZCU104 上,ANN 推理性能为 691.2 GFLOPs/s,而 SNN 可达 5529.6 GFLOPs/s(8×)。在 45 nm 工艺下,ANN 乘法运算消耗 4.6 pJ,SNN 加法运算仅需 0.9 pJ(5.1× 节能)。

这个 energy 和 latency 都是怎么测出来的?在神经模态芯片上部署又是什么硬件?

image.png

4.2. Ablation Study

image.png

注意力分数函数

将提出的 汉明注意力 换成 点积注意力,两项任务的性能均大幅下降:姿态跟踪下降约 15 %,VSS 下降约 10 %。功耗虽略有降低,但原因在于点积的掩蔽效应减少了注意力后的脉冲率。

时空注意力机制

将 联合时空注意力 替换为 层次 或分解注意力,两项任务均出现明显性能退化,主要因为无法充分捕获全局时空特征。这与视频生成领域的最新发现一致。

仅使用 神经元级时间编码 (Neuron-Only,不在时间域做注意力)时,性能有所下降但仍可接受;若完全去掉时间编码(Spatial-Only,每个时间步重置膜电位),性能急剧下滑,说明 SNN 内在的时间编码能力至关重要。

模型初始化

使用 ImageNet-1K 预训练权重 可显著提升下游效果:PA-MPJPE 改善 26 %+,mIoU 提升 19 %。

阈值缩放 s

先前工作固定 s=1/8。如表 6 所示,固定 s 的表现普遍劣于本文提出的s=12Ds=\frac{1}{2D}其中 D 随层级变化,可自适应不同通道宽度,效果更佳。

5. Conclusion

本文提出了 SpikeVideoFormer——一款高效的脉冲驱动视频 Transformer,专为在视频视觉任务中充分发挥 SNN 优势而设计。通过引入 汉明注意力 与 联合时空注意力,SpikeVideoFormer 在视频分类、人体姿态跟踪和视频语义分割三大任务上刷新了现有 SNN 的最佳成绩,并在保持精度的同时,相比最新 ANN 方法显著提升了能效。此外,其线性时间复杂度 O (T) 进一步展示了该 SNN-Transformer 在大规模视频处理中的潜力。未来,我们计划将 SpikeVideoFormer 扩展为更大规模的脉冲骨干网络,支持更广泛的视频任务,如视频理解与生成等。