🤖 machine learning

SAGE: Surrogate-gradient Adaptation via Attention-Guided Entropy for Spiking Transformers

本文介绍了 SAGE，一种用于脉冲神经网络 Transformer 的不确定性调制代理梯度机制，该机制通过利用自注意力熵在训练期间动态调整梯度斜率，从而在不改变推理模型的情况下，实现相对于固定代理基准的一致性准确率提升。

原作者： Kiran Nair, Rodrigue Rizk, KC Santosh

发布于 2026-08-17

📖 1 分钟阅读☕ 轻松阅读

原作者： Kiran Nair, Rodrigue Rizk, KC Santosh

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一个这样的世界：计算机不再是以一种稳定、低鸣的节奏进行数字运算，而是像一座繁忙的萤火虫之城那样进行交流。在这个世界里，信息不再是持续的数据流，而是一系列快速、尖锐的闪烁——要么发生，要么不发生的火花。这就是脉冲神经网络（Spiking Neural Networks, SNNs）的领域，一种受我们自身大脑工作方式启发的类人工智能。这些网络并不使用沉重且耗能的计算，而是仅在必要时发送微小的、二进制的“脉冲”（就像 1 或 0）。这有点像一个社区，你只有在有急事要说时才会去敲门，而不是每隔五分钟就按一次门铃。这使得它们极其高效，非常适合那些需要快速思考却又不想耗尽电池的电池驱动设备或机器人。

然而，教导这些“萤火虫大脑”学习是一件棘手的事情。在传统的 AI 中，我们可以使用平滑的数学方法来确定如何改进网络的预测。但在脉冲网络中，“敲门”是一个全或无（all-or-nothing）的事件；你无法朝着“敲门”迈出一小步，你要么做了，要么没做。这导致学习的数学逻辑失效了。为了解决这个问题，科学家们使用了一个被称为“代理梯度”（surrogate gradient）的巧妙技巧。把它想象成一个训练假人或一个影子，它模仿那次“敲门”，以便老师能够弄清楚如何调整网络。多年来，无论大脑实际在看什么，大家都在为每一个部分使用相同的“影子”。但如果大脑的某些部分感到困惑，需要与那些已经确信的部分不同种类的帮助，该怎么办呢？这就是这篇论文试图解决的谜题。

于是，SAGE（基于注意力引导熵的代理梯度自适应，Surrogate-gradient Adaptation via Attention-Guided Entropy）登场了，这是一种旨在让这些脉冲网络学得更聪明、而非更辛苦的新方法。SAGE 背后的研究人员注意到，在被称为“Transformer”的现代 AI 模型中，网络会以不同的方式关注图像的不同部分。有时，网络对它所看到的东西非常自信（比如一张清晰的猫脸），有时则完全迷失了方向（比如模糊的背景）。旧训练方法的问题在于，它将自信的部分和困惑的部分视为完全相同，使用一个僵化、不变的“影子”来引导学习。

SAGE 改变了游戏规则，它扮演着一位观察团队专注力的睿智教练。它观察网络的注意力是如何分布在不同的“头”（head，可以理解为观察同一场景的不同侦察兵）之中的。如果侦察兵们都看向同一个方向，网络就是自信的；如果他们看向四面八方，网络就是不确定的。SAGE 利用这种“不确定性信号”来实时调整训练中的“影子”。当网络感到困惑时，SAGE 会让训练信号变得更宽泛、更具探索性，允许网络尝试新的事物。当网络感到自信时，SAGE 则会收紧信号，进行精确且细致的调整。

最棒的部分是什么？这种魔法只发生在网络在“教室”里的时候（即训练期间）。一旦网络完成学习并准备好走向现实世界（即推理期间），SAGE 就会消失。网络的运行方式与之前完全一样，没有任何额外的能量消耗，也不会降低速度。在 CIFAR-10 和 CIFAR-100 等图像数据集的测试中，这种方法帮助网络提升了识别图片的能力，与旧有的僵化方法相比，准确率提高了约 1% 到 2%。这是一个微小但意义深远的胜利，表明通过倾听网络自身的困惑，我们可以教导这些高效节能的“萤火虫大脑”更清晰地观察世界。

技术摘要：SAGE —— 基于注意力引导熵的代理梯度自适应脉冲 Transformer

1. 问题陈述

脉冲神经网络（SNNs）通过利用稀疏、事件驱动的计算，为传统的深度神经网络提供了一种高能效的替代方案。然而，训练深层脉冲视觉 Transformer（ViTs）仍然具有挑战性，这是由于脉冲生成函数（Heaviside 阶跃函数）的不可微性导致的。这使得在反向传播过程中必须使用代理梯度（surrogate gradients）来近似导数。

目前的尖端方法通常在所有空间 Token 和训练阶段采用固定或全局统一的代理梯度函数（例如静态 arctangent 或 Fast Sigmoid）。本文认为，这种空间统一的近似在深层脉冲 Transformer 中产生了一个关键的失效模式：

忽略了 Token 级的动态特性： 视觉 Token 表现出不同程度的语义不确定性。高置信度 Token（例如清晰的物体边界）需要精确的梯度更新，而高不确定性 Token（例如模糊的背景）则受益于更宽泛、更具探索性的梯度流。
优化失败： 在所有 Token 上强制执行静态代理窗口，无法提供上下文信用分配（contextual credit assignment）。在深层网络中，这会导致严重的梯度消失，并导致注意力头塌陷为“死头”（dead heads，即膜电位保持在阈值以下，导致权重更新停滞）。

核心挑战在于开发一种自适应代理机制，该机制能够在不损害推理所需的确定性、低功耗二进制执行的前提下，动态调整局部 Token 不确定性的梯度流。

2. 方法论：SAGE 框架

作者提出了 SAGE（基于注意力引导熵的代理梯度自适应），这是一个通过分析注意力熵的离散度来调节训练期间代理梯度斜率的不确定性感知框架。

2.1 核心概念

SAGE 仅在训练阶段运行。它通过分析自注意力头中注意力熵的离散度来估计每个 Transformer 块的不确定性。该信号用于自适应地调整代理梯度函数的斜率（ $\alpha$ ）。至关重要的是，推理模型保持不变；不确定性控制器在训练后被移除。

2.2 注意力不确定性估计

该框架利用了 Spikformer 架构的多头自注意力（MHA）模块：

原始注意力分数： 标准的 Spikformer 计算原始分数 $S = s QK^\top$ ，且不进行 Softmax 归一化。
辅助分布： 仅用于不确定性估计，SAGE 通过应用 detach 操作和带有固定温度（ $T_{ent} = 0.25$ ）的 Softmax 来构建一个临时概率分布 $A$ ：
$A = \text{Softmax}\left(\frac{\text{detach}(S)}{T_{ent}}\right)$
熵计算： 为每个注意力头 $i$ 计算归一化香农熵 $E_i$ 。
不确定性信号： SAGE 并非使用平均熵，而是计算所有 $H$ 个头之间熵的标准差（离散度）：
$D = \text{Std}(E_1, E_2, \dots, E_H)$
高离散度表示头之间存在分歧（高不确定性），而低离散度则意味着注意力模式具有一致性。

2.3 不确定性引导的代理调制

离散度统计量 $D$ 被处理以生成块特定的代理斜率 $\alpha$ ：

平滑： 使用指数移动平均（EMA）对瞬时离散度进行平滑，以确保稳定性。
归一化： 使用运行均值和标准差统计量对平滑后的值进行归一化，然后在所有 Transformer 块之间进行中心化，从而获得相对不确定性得分 $\tilde{z}$ 。
自适应映射： 死区控制器抑制微小波动。通过双曲正切函数对斜率进行调制：
$\hat{\alpha} = \begin{cases} 4.0 & \text{如果 } |\tilde{z}| < 0.25 \\ 4.0 + 0.5 \tanh(\tilde{z}) & \text{否则} \end{cases}$
最终斜率被限制在 $[3, 5]$ 区间内。
集成： 该自适应 $\alpha$ 仅在反向传播期间应用。前向传播和推理架构与基准模型完全相同。

2.4 计算开销

增加的计算涉及计算熵（$O(HN) $）和标准差（$ O(H) $），与注意力机制的$ O(N^2)$ 复杂度相比可以忽略不计。性能分析显示，在训练期间，每个 mini-batch 的额外开销约为 0.03 ms，而在推理期间开销为零。

3. 主要贡献

论文概述了三个主要贡献：

首个不确定性感知框架： 引入了第一个专门为脉冲 Transformer 设计的、利用注意力衍生不确定性的自适应代理梯度框架。
熵离散度作为信号： 证明了注意力熵在各头之间的离散度（标准差）是比其他指标（如 Gini 不纯度或平均熵）更具判别力和可靠性的在线不确定性信号，用于代理自适应。
性能与效率： 在竞争性数据集上进行了验证，显示出相对于固定和可学习基准的一致性改进，且具有极低的训练开销且无推理成本。

4. 实验结果

实验在 CIFAR-10、CIFAR-100 和 ImageNet-200 上使用 Spikformer-4-384 骨干网络进行。

准确率增益： SAGE 在 CIFAR 数据集上较固定代理基准实现了持续的改进。
- CIFAR-10： 96.32% Top-1 准确率（基准 Spikformer 为 95.51%）。
- CIFAR-100： 78.69% Top-1 准确率（基准为 78.21%）。
- ImageNet-200： 在此基准上，SAGE 实现了 60.60% Top-1 准确率，这与基准 Spikformer 的 70.24% 相比有所下降（表 2）。论文指出，虽然 SAGE 在 CIFAR 数据集上提高了性能，但在 ImageNet-200 上的结果低于基准，但该方法仍因其在低延迟机制下的有效性以及能够保留原始架构而被强调。
时间鲁棒性： SAGE 在较低的模拟时间步（ $T \le 4$ ）下表现出最显著的改进，特别是在 $T=4$ 的标准设置下。随着时间步增加（ $T=8, 12$ ），性能差距缩小，这表明该方法对于低延迟应用最为有效。
对比： SAGE 在多个数据集和时间步长下均优于固定代理梯度和“可学习”配置（其中一个全局参数被训练）。

5. 重要性与主张

论文将 SAGE 定位为一种 即插即用优化策略，解决了深层脉冲 Transformer 中静态代理梯度的特定局限性。

架构保留： 该方法保留了原始的网络架构、模型参数和推理流水线。它没有为模型本身引入额外的可学习参数，也不会增加推理延迟。
上下文信用分配： 通过根据局部 Token 不确定性动态调整梯度斜率，SAGE 缓解了深层网络中的“死头”现象和梯度消失问题，使网络在训练期间能更好地处理模糊特征。
实用性： 作者声称 SAGE 是一种轻量级解决方案，可以无缝集成到现有的训练流水线中，无需更改部署方式，使其能够立即应用于广泛的 SNN 及基于 Transformer 的 SNN 模型。

作者得出结论，不确定性引导的代理优化代表了改进深层 SNN 训练的一个有效方向，特别是对于那些要求低延迟执行的应用。