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

1. Problem Statement

Spiking Neural Networks (SNNs) offer an energy-efficient alternative to conventional deep neural networks by utilizing sparse, event-driven computation. However, training deep Spiking Vision Transformers (ViTs) remains challenging due to the non-differentiable nature of the spike generation function (Heaviside step function). This necessitates the use of surrogate gradients during backpropagation to approximate derivatives.

Current state-of-the-art approaches typically employ fixed or globally uniform surrogate gradient functions (e.g., static arctangent or Fast Sigmoid) across all spatial tokens and training stages. The paper argues that this spatially uniform approximation creates a critical failure mode in deep Spiking Transformers:

• Token-Level Dynamics Ignored: Visual tokens exhibit varying levels of semantic uncertainty. High-confidence tokens (e.g., distinct object boundaries) require precise gradient updates, while high-uncertainty tokens (e.g., ambiguous backgrounds) benefit from broader, exploratory gradient flow.
• Optimization Failure: Enforcing a static surrogate window across all tokens fails to provide contextual credit assignment. In deep layers, this leads to severe gradient vanishing and causes attention heads to collapse into "dead heads" (where membrane potentials remain subthreshold and weight updates stall).

The core challenge is developing an adaptive surrogate mechanism that dynamically tunes gradient flow to local token uncertainty without compromising the deterministic, low-power binary execution required for inference.

2. Methodology: SAGE Framework

The authors propose SAGE (Surrogate-gradient Adaptation via Attention-Guided Entropy), an uncertainty-aware framework that modulates the surrogate-gradient slope during training based on the dispersion of attention entropy.

2.1 Core Concept

SAGE operates exclusively during the training phase. It estimates the uncertainty of each transformer block by analyzing the dispersion of attention entropy across self-attention heads. This signal is used to adaptively adjust the slope ($\alpha$) of the surrogate gradient function. Crucially, the inference model remains unchanged; the uncertainty controller is removed post-training.

2.2 Attention Uncertainty Estimation

The framework leverages the Spikformer architecture's Multi-Head Self-Attention (MHA) module:

1. Raw Attention Scores: The standard Spikformer computes raw scores $S = s QK^\top$ without softmax normalization.
2. Auxiliary Distribution: For uncertainty estimation only, SAGE constructs a temporary probability distribution $A$ by applying a detach operation and a Softmax with a fixed temperature ($T_{ent} = 0.25$):
   $$A = \text{Softmax}\left(\frac{\text{detach}(S)}{T_{ent}}\right)$$
3. Entropy Calculation: The normalized Shannon entropy $E_i$ is computed for each attention head $i$.
4. Uncertainty Signal: Instead of using the mean entropy, SAGE calculates the standard deviation (dispersion) of the entropies across all $H$ heads:
   $$D = \text{Std}(E_1, E_2, \dots, E_H)$$
   A high dispersion indicates disagreement among heads (high uncertainty), while low dispersion implies consistent attention patterns.

2.3 Uncertainty-Guided Surrogate Modulation

The dispersion statistic $D$ is processed to generate a block-specific surrogate slope $\alpha$:

1. Smoothing: The instantaneous dispersion is smoothed using an Exponential Moving Average (EMA) to ensure stability.
2. Normalization: The smoothed value is normalized using running mean and standard deviation statistics, then centered across all transformer blocks to obtain a relative uncertainty score $\tilde{z}$.
3. Adaptive Mapping: A dead-zone controller suppresses minor fluctuations. The slope is modulated via a hyperbolic tangent function:
   $$\hat{\alpha} = \begin{cases} 4.0 & \text{if } |\tilde{z}| < 0.25 \\ 4.0 + 0.5 \tanh(\tilde{z}) & \text{otherwise} \end{cases}$$
   The final slope is clipped to the interval $[3, 5]$.
4. Integration: This adaptive $\alpha$ is applied only during the backward pass. The forward pass and inference architecture remain identical to the baseline Spikformer.

2.4 Computational Overhead

The additional computation involves calculating entropy ($O(HN)$) and standard deviation ($O(H)$), which is negligible compared to the $O(N^2)$ complexity of the attention mechanism. Profiling shows an overhead of approximately 0.03 ms per mini-batch during training and zero overhead during inference.

3. Key Contributions

The paper outlines three primary contributions:

1. First Uncertainty-Aware Framework: Introduction of the first adaptive surrogate-gradient framework specifically designed for spiking transformers that utilizes attention-derived uncertainty.
2. Entropy Dispersion as a Signal: Demonstration that the dispersion (standard deviation) of attention entropy across heads is a more discriminative and reliable online uncertainty signal for surrogate adaptation than alternative metrics (such as Gini impurity or mean entropy).
3. Performance and Efficiency: Validation on competitive datasets showing consistent improvements over fixed and learnable baselines with negligible training overhead and no inference cost.

4. Experimental Results

Experiments were conducted on CIFAR-10, CIFAR-100, and ImageNet-200 using the Spikformer-4-384 backbone.

• Accuracy Gains: SAGE achieved consistent improvements over fixed-surrogate baselines on CIFAR datasets.
  • CIFAR-10: 96.32% Top-1 accuracy (vs. 95.51% for baseline Spikformer).
  • CIFAR-100: 78.69% Top-1 accuracy (vs. 78.21% for baseline).
  • ImageNet-200: On this benchmark, SAGE achieved 60.60% Top-1 accuracy, which represents a decrease compared to the baseline Spikformer's 70.24% (Table 2). The paper notes that while SAGE improves performance on CIFAR datasets, the results on ImageNet-200 were lower than the baseline, though the method is highlighted for its effectiveness in low-latency regimes and its ability to preserve the original architecture.
• Temporal Robustness: SAGE showed the most pronounced improvements at lower simulation time steps ($T \le 4$), particularly at the standard setting of $T=4$. As time steps increased ($T=8, 12$), the performance gap narrowed, suggesting the method is most effective for low-latency applications.
• Comparison: SAGE outperformed both fixed surrogate gradients and a "Learnable" configuration (where a single global parameter is trained) across multiple datasets and time steps.

5. Significance and Claims

The paper positions SAGE as a plug-and-play optimization strategy that addresses the specific limitations of static surrogate gradients in deep Spiking Transformers.

• Architectural Preservation: The method preserves the original network architecture, model parameters, and inference pipeline. It introduces no additional learnable parameters for the model itself and incurs no inference latency.
• Contextual Credit Assignment: By dynamically tuning the gradient slope based on local token uncertainty, SAGE mitigates the "dead head" phenomenon and gradient vanishing in deep layers, allowing the network to better handle ambiguous features during training.
• Practicality: The authors claim that SAGE is a lightweight solution that can be seamlessly integrated into existing training pipelines, making it immediately applicable to a broad range of SNN and Transformer-based SNN models without requiring deployment changes.

The authors conclude that uncertainty-guided surrogate optimization represents an effective direction for improving the training of deep SNNs, particularly for applications requiring low-latency execution.