SAGE: 基于注意力引导熵的替代梯度自适应用于脉冲Transformer

# SAGE：基于注意力引导熵的脉冲Transformer自适应替代梯度方法
来源：https://arxiv.org/html/2608.13702
Kiran Nair谢谢：通讯作者：kiran\.prasannannair@coyotes\.usd\.edu所属机构：USD人工智能研究实验室所属机构：计算机科学系所属机构：南达科他大学所属机构：美国南达科他州维米利恩市 57069
Rodrigue Rizk所属机构：USD人工智能研究实验室所属机构：计算机科学系所属机构：南达科他大学所属机构：美国南达科他州维米利恩市 57069
KC Santosh所属机构：USD人工智能研究实验室所属机构：计算机科学系所属机构：南达科他大学所属机构：美国南达科他州维米利恩市 57069

###### 摘要

脉冲神经网络（SNN）通过利用稀疏的事件驱动计算，为传统深度神经网络提供了一种高能效替代方案，但其训练仍具挑战性，因为非可微的脉冲函数需要依赖固定形状的替代梯度，而该形状在不同层和训练阶段可能并非最优。本文提出SAGE——一种面向脉冲Transformer的不确定性调制替代梯度机制。SAGE通过归一化自注意力熵估计块级不确定性，并利用该信号在训练过程中自适应调整替代梯度斜率，同时保持推理模型不变。通过仅调制训练时的替代参数，所提方法在保持原始架构和部署成本的同时，提升了优化灵活性。在CIFAR-10/100上的实验表明，SAGE相较于固定替代梯度基线实现了精度提升，在多个仿真时间步上持续获得1-2%的增益。这些结果突显了基于注意力的不确定性作为轻量级训练信号，用于脉冲Transformer中自适应替代梯度学习的潜力。

## 1 引言

脉冲神经网络（SNN）作为第三代神经网络模型，已成为深度学习中一种高能效、事件驱动的范式24 (https://arxiv.org/html/2608.13702#bib.bib33); 32 (https://arxiv.org/html/2608.13702#bib.bib35)。SNN用离散二值脉冲（00或11）替代连续浮点激活，在神经形态硬件上运行时，将资源密集的乘加操作替换为稀疏的、仅加法（AC）操作，显著降低了计算能耗14 (https://arxiv.org/html/2608.13702#bib.bib40); 19 (https://arxiv.org/html/2608.13702#bib.bib34)。基于这些生物学上合理的特性，近期架构进展已成功地将自注意力机制引入SNN。其中最著名的是Spikformer架构52 (https://arxiv.org/html/2608.13702#bib.bib10)，它建立了脉冲自注意力（SSA），通过消除softmax归一化来跨脉冲形式的Query、Key和Value张量计算稀疏的AC注意力图。这一突破表明，脉冲视觉Transformer（ViTs）可以在保持超低硬件能耗的同时，在大规模视觉基准上实现具有竞争力的准确率。

尽管前向计算高效，但训练深度Spiking ViTs仍是一个根本性挑战。由于Heaviside脉冲生成函数不可微，无法直接应用时间反向传播（BPTT）。因此，现代SNN训练依赖于在反向传播过程中近似导数的替代梯度21 (https://arxiv.org/html/2608.13702#bib.bib3); 27 (https://arxiv.org/html/2608.13702#bib.bib5); 43 (https://arxiv.org/html/2608.13702#bib.bib36); 50 (https://arxiv.org/html/2608.13702#bib.bib41)。然而，传统的替代梯度函数，如静态反正切、快速Sigmoid或分段线性导数，对所有空间序列维度应用固定或全局统一的梯度响应窗口2 (https://arxiv.org/html/2608.13702#bib.bib38); 51 (https://arxiv.org/html/2608.13702#bib.bib39)。尽管近期工作探索了时间上可学习的替代梯度以跟踪时间步动态51 (https://arxiv.org/html/2608.13702#bib.bib39)，但现有范式在反向传播过程中对所有空间token一视同仁，忽略了视觉Transformer层中固有的复杂、token级的特征动态。

这种空间均匀的梯度近似在深度Spiking Transformer中导致了一个关键故障模式。视觉token表现出截然不同的语义不确定性水平：高置信度token（如清晰的物体边界）需要窄而精确的梯度更新，而高不确定性token（如模糊的背景杂乱或复杂纹理）则需要更宽、更具探索性的梯度流来发现最优表示。通过在所有token上强制使用静态的替代窗口，标准反向传播无法提供上下文相关的信用分配。在深层网络中，这种token盲的梯度估计加剧了严重的梯度消失，并导致注意力头坍缩为“死头”，即膜电位保持在阈下，权重更新完全停滞52 (https://arxiv.org/html/2608.13702#bib.bib10); 51 (https://arxiv.org/html/2608.13702#bib.bib39)。因此，扩展Spiking ViTs的一个主要挑战在于开发一种自适应替代机制，能够在不影响前向传播确定性、低功耗二值执行的情况下，根据局部token不确定性动态调整梯度流。

受此挑战启发，我们提出了SAGE（基于注意力引导熵的替代梯度自适应），一个面向Spiking ViTs的不确定性感知替代梯度框架。SAGE不在整个训练过程中使用固定或全局可学习的替代梯度，而是从自注意力头的注意力熵离散度中估计每个Transformer块的不确定性，并在反向传播过程中自适应地调制替代梯度斜率。所提框架仅在训练期间运行，无需架构修改、额外的推理时参数或更改前向计算。在CIFAR-10、CIFAR-100和ImageNet-200上的大量实验表明，SAGE在保持原始Spikformer效率和部署特性的同时，持续提升了训练效果。具体贡献如下：

1. 1\. 我们首次提出了面向脉冲Transformer的不确定性感知自适应替代梯度框架。
2. 2\. 我们证明，跨注意力头的熵离散度为替代梯度自适应提供了可靠的在线不确定性信号，并且在表征Transformer块不确定性方面，其区分度显著高于其他度量指标。
3. 3\. 我们在有竞争力的数据集上验证了SAGE，它持续优于固定和可学习的替代梯度基线，Top-1准确率提升高达1-2%，同时每个小批次仅增加0\.03毫秒的训练开销。

本文其余部分组织如下。第2节 (https://arxiv.org/html/2608.13702#S2) 介绍了所提框架的动机，通过理论见解和支持性文献确立了对不确定性感知替代梯度自适应的需求。第3节 (https://arxiv.org/html/2608.13702#S3) 回顾了替代梯度学习和Spiking ViTs的现有文献。第4节 (https://arxiv.org/html/2608.13702#S4) 介绍了所提的SAGE框架，详述了不确定性估计策略及其在Spikformer训练中的集成。第5节 (https://arxiv.org/html/2608.13702#S5) 描述了实验设置并展示了全面的评估。第6节 (https://arxiv.org/html/2608.13702#S6) 对实验结果进行了详细分析和讨论，而第7节 (https://arxiv.org/html/2608.13702#S7) 则对本文进行了总结。

## 2 动机

深度SNN使用替代梯度进行训练，这些梯度在反向传播过程中近似不可微脉冲函数的导数。在过去十年中，已提出了众多替代梯度公式，主要差异在于其函数形式49 (https://arxiv.org/html/2608.13702#bib.bib42)、平滑度特性35 (https://arxiv.org/html/2608.13702#bib.bib43)和梯度缩放策略43 (https://arxiv.org/html/2608.13702#bib.bib36)。尽管存在这些差异，现有方法对每个神经元采用固定的替代响应，无论底层表示的置信度如何，都一视同仁地处理所有脉冲事件。这一假设与现代学习范式形成对比，在后者中，中间表示对下游预测的贡献并不相等40 (https://arxiv.org/html/2608.13702#bib.bib32)，并且不确定性估计被广泛用于识别模糊或不可靠的特征10 (https://arxiv.org/html/2608.13702#bib.bib44)。因此，将相同的替代梯度应用于编码高确定性和高不确定性信息的神经元，可能导致优化过程中的信用分配次优27 (https://arxiv.org/html/2608.13702#bib.bib5)。

这些观察自然引出了以下问题：当底层表示表现出不同程度的不确定性时，替代梯度是否仍应保持均匀？我们假设，替代梯度调制应反映编码表示的不确定性。直观地说，高置信度表示需要较少的替代平滑，而高不确定性表示可能受益于更宽的替代支持，以维持有效的梯度传播。这一直觉与现代深度网络中自适应计算和不确定性感知学习的广泛原理一致40 (https://arxiv.org/html/2608.13702#bib.bib32); 10 (https://arxiv.org/html/2608.13702#bib.bib44)。受此观察启发，我们提出了一种不确定性调制的替代梯度，它利用注意力导出的不确定性动态调整后向替代响应，同时保留Spiking ViTs的确定性前向动态。所提公式在命题1 (https://arxiv.org/html/2608.13702#Thmproposition1) 中进行了理论分析，该命题建立了不确定性与替代梯度幅度之间的单调关系。

###### 命题 1\.

设 \(z_c\) 表示一个Transformer块的中心归一化熵离散度统计量。自适应替代梯度斜率定义为
\(\alpha(z_c) = \begin{cases} \alpha_0, & \|z_c\| < \delta, \\ \alpha_0 + \beta \tanh(z_c), & \|z_c\| \geq \delta, \end{cases}\)  (1)
其中 \(\alpha_0 > 0\)，\(\beta > 0\)，且 \(\delta > 0\)。那么，在死区之外，\(\alpha(z_c)\) 是 \(z_c\) 的单调递增函数。

###### 证明\.

对于 \(\|z_c\| \geq \delta\)，
\(\frac{d\alpha}{dz_c} = \beta (1 - \tanh^2(z_c))\).  (2)
由于 \(\beta > 0\) 且对于所有有限的 \(z_c\)，\(1 - \tanh^2(z_c) > 0\)，
\(\frac{d\alpha}{dz_c} > 0\).  (3)
因此，在死区之外，\(\alpha\) 相对于中心不确定性统计量是严格单调的。在死区内，\(\alpha = \alpha_0\) 是常数，有意抑制小的不确定性波动。 ∎

## 3 相关工作

本节回顾与所提SAGE框架最密切相关的三个研究方向。我们首先总结脉冲神经网络直接训练的进展，其中替代梯度已成为优化深度SNN的标准方法。然后，我们讨论Spiking ViTs的演变，重点介绍基于自注意力机制的最新架构发展。最后，我们回顾现有的替代梯度优化策略及其局限性，从而为脉冲Transformer中需要不确定性引导的自适应替代框架提供了动机。

##### 脉冲神经网络与直接训练

SNN代表了一种受生物启发的事件驱动范式，通过传输离散二值脉冲而非连续浮点激活来实现显著的能效11 (https://arxiv.org/html/2608.13702#bib.bib1); 38 (https://arxiv.org/html/2608.13702#bib.bib2)。扩展SNN的早期工作依赖于人工神经网络（ANN）到SNN的转换，但这通常会导致过高的延迟和较长的推理时间步13 (https://arxiv.org/html/2608.13702#bib.bib37)。为克服这些延迟瓶颈，通过时空反向传播（STBP）和BPTT进行直接训练成为了优化深度架构的主流框架21 (https://arxiv.org/html/2608.13702#bib.bib3); 43 (https://arxiv.org/html/2608.13702#bib.bib36)。直接训练算法在反向传播过程中用平滑的替代梯度替代不可微的Heaviside阶跃函数，从而实现具有竞争力的低延迟执行27 (https://arxiv.org/html/2608.13702#bib.bib5)。最新进展已将直接SNN训练扩展到多个领域，包括动态脉冲图网络46 (https://arxiv.org/html/2608.13702#bib.bib6)、异步事件视觉处理器48 (https://arxiv.org/html/2608.13702#bib.bib7)和注意力引导的脉冲架构45 (https://arxiv.org/html/2608.13702#bib.bib4)。然而，一个根本局限性依然存在：直接训练算法严重依赖静态或全局统一的替代导数函数27 (https://arxiv.org/html/2608.13702#bib.bib5); 51 (https://arxiv.org/html/2608.13702#bib.bib39)，无法适应复杂特征表示中的局部空间变化。

##### 脉冲视觉Transformer

为结合自注意力40 (https://arxiv.org/html/2608.13702#bib.bib32) 的表征能力与神经形态处理的低功耗特性，近期研究已将ViTs调整适应脉冲域操作。开创性的Spikformer架构52 (https://arxiv.org/html/2608.13702#bib.bib10) 引入了SSA，它消除了softmax归一化，以跨脉冲形式的Query、Key和Value张量计算稀疏的AC注意力图。在此基础上，后续工作引入了多尺度脉冲ViTs47 (https://arxiv.org/html/2608.13702#bib.bib11)、混合事件检测模型44 (https://arxiv.org/html/2608.13702#bib.bib12)、类似扫视的生物视觉机制41 (https://arxiv.org/html/2608.13702#bib.bib9)，以及ResNets与ViTs之间的结构桥梁34 (https://arxiv.org/html/2608.13702#bib.bib8)。此外，还探索了门控机制26 (https://arxiv.org/html/2608.13702#bib.bib14) 和动态时间步分配框架7 (https://arxiv.org/html/2608.13702#bib.bib13)，以增强信息流控制和降低延迟。然而，尽管这些架构优化了前向传播执行和结构特征路由，但它们的反向优化仍受限于静态替代函数52 (https://arxiv.org/html/2608.13702#bib.bib10)，使得深度脉冲Transformer层容易出现梯度消失和死注意力头。

##### 替代梯度优化

克服脉冲激活的不可微性激发了大量关于替代梯度公式化的研究。传统的直接训练范式采用固定解析导数，如快速Sigmoid、反正切或分段线性函数2 (https://arxiv.org/html/2608.13702#bib.bib38)。为缓解梯度消失和不匹配，研究人员开发了自适应和可学习的替代梯度方法。