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

Kiran Nair; Rodrigue Rizk; KC Santosh

Source-linked AI summary

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

Kiran Nair, Rodrigue Rizk, KC Santosh

arXiv:2608.13702v1Aug 13, 2026cs.LGcs.AIcs.CVcs.NE

TL;DR

Training spiking transformers is hindered by fixed surrogate gradients that may not suit representations with different uncertainty levels. SAGE adapts the training-time surrogate slope using attention-derived uncertainty, improving accuracy by up to 1–2% over fixed and learnable baselines without increasing inference cost.

Problem
Fixed surrogate gradients provide uniform responses despite differing representation uncertainty, limiting contextual credit assignment during spiking-transformer training.
Method
SAGE estimates uncertainty from attention-entropy dispersion across self-attention heads and adaptively modulates the surrogate-gradient slope during training without architectural changes.
Results
Up to 1–2% Top-1 accuracy gains over fixed and learnable surrogate baselines were observed across evaluated datasets and simulation settings.
Takeaways & Limitations
SAGE provides a lightweight, plug-and-play optimization strategy for spiking transformers with no additional inference cost.
Takeaways & Limitations
The study focuses on attention-based spiking transformers, so applying SAGE to non-transformer architectures requires alternative uncertainty measures.

Abstract

from arXiv · show

Spiking neural networks (SNNs) offer an energy-efficient alternative to conventional deep neural networks by exploiting sparse event-driven computation, but their training remains challenging because the non-differentiable spike function requires surrogate gradients whose fixed shape may be suboptimal across layers and training stages. In this work, we introduce SAGE, an uncertainty-modulated surrogate-gradient mechanism for Transformer-based SNNs. SAGE estimates block-level uncertainty from normalized self-attention entropy and uses this signal to adapt the surrogate-gradient slope during training while leaving the inference model unchanged. By modulating only the training-time surrogate parameter, the proposed method preserves the original architecture and deployment cost while improving optimization flexibility. Experiments on CIFAR-10/100 demonstrate that SAGE achieves improved accuracy over fixed-surrogate baselines, with results up to 1-2\% consistent gains across multiple simulation time steps. These results highlight the potential of attention-derived uncertainty as a lightweight training signal for adaptive surrogate-gradient learning in transformer-based SNNs.

1 Introduction

SAGE addresses the training difficulty of Spiking Vision Transformers caused by non-differentiable spikes and static surrogate gradients. It adapts the surrogate-gradient slope using attention-entropy uncertainty, improving accuracy by up to 1–2% with 0.03 ms training overhead per mini-batch.

Background: SNNs use discrete binary spikes and sparse addition-only operations, enabling energy-efficient event-driven computation on neuromorphic hardware.Spiking self-attention eliminates softmax normalizations while retaining competitive accuracy and ultra-low hardware energy profiles.
Problem: Deep Spiking ViT training is difficult because the non-differentiable Heaviside function prevents direct BPTT, requiring surrogate gradients for backward propagation.Conventional surrogate gradients may impose an unsuitable static approximation across layers and training conditions.
Problem: Spatially uniform surrogate gradients mismatch token uncertainty: confident tokens need narrow updates, whereas ambiguous tokens require broader exploratory gradient flow.This mismatch creates a critical failure mode in deep Spiking transformers.
SAGE Framework: SAGE estimates transformer-block uncertainty from attention-entropy dispersion across self-attention heads and adaptively modulates surrogate-gradient slope during backpropagation.The framework is uncertainty-aware and operates during training rather than changing the deployed model architecture.
Results: Up to 1-2% Top-1 accuracy gains and 0.03 ms training overhead per mini-batch distinguish SAGE from fixed and learnable surrogate baselines.Entropy dispersion is reported as a reliable online uncertainty signal and more discriminative than alternative transformer-block uncertainty metrics.

2 Motivation

Existing surrogate-gradient methods use a fixed response across neurons despite differing representation uncertainty. SAGE motivates uncertainty-dependent modulation, with less smoothing for confident representations and broader support for uncertain ones.

Motivation: Existing surrogate-gradient formulations differ in functional form, smoothness, and gradient scaling, yet use a fixed surrogate response for every neuron.These methods approximate the derivative of the non-differentiable spike function during backpropagation.
Motivation: The paper asks whether surrogate gradients should remain uniform when encoded representations exhibit different uncertainty levels.This question motivates adapting the surrogate response to representation uncertainty.
Motivation: The proposed hypothesis is to modulate surrogate gradients according to the uncertainty of the encoded representation.The approach is presented as an alternative to uniform surrogate-gradient treatment.
Motivation: Confident representations are expected to require less surrogate smoothing, whereas uncertain representations may benefit from broader surrogate support for effective gradient propagation.This contrast provides the intuition for uncertainty-dependent surrogate-gradient modulation.
Motivation: Outside the dead-zone region, α(z_c) increases monotonically with the centered uncertainty statistic, while within the dead zone α = α_0 remains constant.The constant dead-zone behavior intentionally suppresses small uncertainty fluctuations.

3 Related Works

This section reviews three research directions related to SAGE: direct training of spiking neural networks, spiking vision transformers, and surrogate-gradient optimization. It emphasizes advances, architectural developments, and limitations motivating the proposed framework.

3 Related Works: The related-work discussion covers direct SNN training, spiking Vision Transformers, and surrogate-gradient optimization.These directions correspond to the section’s three stated research areas.
Spiking Neural Networks & Direct Training: SNNs use event-driven binary spikes for energy efficiency, while ANN-to-SNN conversion can incur excessive latency and high inference time steps.Direct training with surrogate gradients is presented as an approach for optimizing deep SNNs.
Spiking Vision Transformers: Spiking Vision Transformers adapt self-attention to neuromorphic processing, with Spikformer introducing softmax-free sparse AC attention over spike-form Query, Key, and Value tensors.This work combines self-attention’s representation capability with the low-power characteristics of neuromorphic processing.
Surrogate Gradient Optimization: Surrogate-gradient research addresses nondifferentiable spiking activations through fixed analytical derivatives and adaptive or learnable surrogate functions.Examples of fixed functions include Fast Sigmoid, arctangent, and piecewise linear derivatives.

4 Method

SAGE adapts surrogate-gradient slopes during training using attention-entropy dispersion as a block-level uncertainty signal, while preserving the original Spikformer inference architecture and cost. The method normalizes per-head attention entropy, smooths and centers its dispersion, then applies bounded block-specific modulation around α = 4.

Training-time framework: SAGE operates only during training, adaptively modulating surrogate gradients while leaving the Spikformer forward architecture and inference procedure unchanged.The uncertainty estimation module and adaptive controller are removed during inference, preserving the inference graph, computational complexity, and latency.
Attention-derived uncertainty: SAGE constructs detached, temperature-scaled auxiliary attention distributions for uncertainty estimation without altering the forward attention computation.The entropy temperature is fixed to T_ent = 0.25 based on OFAT sensitivity analysis.
Attention-derived uncertainty: Block uncertainty is the standard deviation of normalized per-head entropies, where larger dispersion indicates greater disagreement and higher uncertainty.Entropy normalization confines each entropy to [0, 1], and the dispersion statistic drives the adaptive controller.
Adaptive surrogate-gradient control: SAGE smooths dispersion with an EMA, normalizes and centers it across transformer blocks, then applies a dead-zone controller before changing surrogate-gradient slopes.These operations produce relative block-wise uncertainty and suppress minor fluctuations around the block-wise mean.
Adaptive surrogate-gradient control: The tanh mapping smoothly modulates slopes around α = 4, with final clipping constraining the surrogate parameter to [3, 5].The centering operation makes adaptation relative to uncertainty across transformer blocks.
Implementation: SAGE adds negligible computation relative to self-attention, estimating per-head entropy in O(HN) and dispersion across heads in O(H).The dominant scaled dot-product attention computation remains O(N^2).

5 Experimental Results

Section 5 evaluates SAGE on CIFAR-10, CIFAR-100, and ImageNet-200 using the Spikformer framework, comparing it with original and recent SNN methods. Experiments also examine adaptive surrogate-gradient configurations, component contributions, and robustness across simulation timesteps.

Benchmark comparisons: Experiments compare SAGE with the original Spikformer baseline, recent state-of-the-art SNNs, and representative Spiking Transformer methods.Reported comparisons include Hybrid Training, DIET-SNN, STBP-tdBN, TET, RMP-SNN, QCFS, and Spikformer.
Surrogate-gradient configurations: The ablation study compares Fixed, Learnable, and SAGE surrogate-gradient configurations.Fixed uses a conventional sigmoid surrogate with constant slope α = 4; Learnable jointly trains a single surrogate parameter; SAGE uses uncertainty-driven entropy modulation.
Datasets: SAGE is evaluated on CIFAR-10, CIFAR-100, and ImageNet-200 image-classification benchmarks.CIFAR-10 and CIFAR-100 each contain 60,000 32×32 RGB images, while CIFAR-100 provides 100 fine-grained classes.
Implementation: SAGE is implemented on the official Spikformer framework without modifying the network architecture or inference pipeline.Unless otherwise specified, experiments use the Spikformer-4-384 backbone with four transformer blocks and four simulation time steps.
Temporal robustness: All three surrogate-gradient configurations are evaluated under different simulation timesteps to investigate robustness across temporal resolutions.SAGE adaptively modulates the training-time slope while leaving the inference model unchanged.

6 Discussion

SAGE delivers competitive or superior accuracy across evaluated datasets while preserving the original Spikformer inference architecture. Its attention-based uncertainty signal reflects block- and head-wise training dynamics, but current applicability is centered on attention-based spiking transformers.

Performance Analysis: SAGE achieves competitive or superior classification accuracy against fixed α = 4 and learnable surrogate formulations across the evaluated datasets.Under the standard Spikformer setting, T = 4, SAGE attains the highest Top-1 accuracy on both CIFAR-10 and CIFAR-100.
Understanding SAGE: Normalized attention entropy reveals clear block-wise and head-wise variations in attention dynamics throughout training.These variations motivated using attention-derived uncertainty to guide SAGE’s adaptive surrogate-gradient mechanism.
Practical Implications: Because SAGE operates exclusively during training, the forward inference graph remains identical to the original Spikformer architecture.The method can be integrated into existing surrogate-gradient training pipelines as a plug-and-play optimization strategy.
Practical Implications: SAGE requires no architectural modifications or deployment changes, supporting application across SNNs and transformer-based SNN models.This preserves the deployment setup while enabling the proposed training-time optimization.
Limitations and Future Work: The current study is limited to attention-based spiking transformers, where uncertainty can be estimated naturally from self-attention distributions.Extending SAGE to convolutional or other non-transformer SNNs will require alternatives such as spike-rate or membrane-potential statistics.

7 Conclusion

SAGE is an uncertainty-aware surrogate-gradient adaptation framework for Spiking ViTs that uses attention-entropy dispersion to adjust training while preserving the original inference setup. Experiments on CIFAR-10, CIFAR-100, and ImageNet-200 showed consistent improvement or competitive classification performance.

7 Conclusion: SAGE presents an uncertainty-aware surrogate-gradient adaptation framework for Spiking ViTs.The framework is introduced as the paper’s central contribution.
7 Conclusion: SAGE dynamically modulates the surrogate-gradient slope during training using attention-entropy dispersion across self-attention heads.The adaptation signal is derived from the dispersion of attention entropy across heads.
7 Conclusion: SAGE preserves the original network architecture, model parameters, and inference pipeline while adapting surrogate gradients.The adaptation occurs during training without changing the stated deployment components.
7 Conclusion: Experiments on CIFAR-10, CIFAR-100, and ImageNet-200 consistently improved or maintained competitive classification performance.The conclusion reports this outcome across all three listed datasets.

A Additional Analysis of Uncertainty Measures

The analysis compares normalized attention entropy with Gini impurity and finds entropy provides a richer uncertainty signal for adapting surrogate gradients. Based on this sensitivity analysis, SAGE uses an entropy temperature of τe = 0.25 throughout experiments.

Uncertainty-measure comparison: Gini impurity remains highly saturated across all four Spikformer blocks, showing only minor temporal fluctuations under all tested attention temperatures.The comparison indicates limited sensitivity to changes in attention uncertainty relative to normalized attention entropy.
Entropy-temperature sensitivity: Lower entropy temperatures preserve substantially greater normalized-attention-entropy dispersion, providing a richer uncertainty signal for SAGE’s adaptive controller.The OFAT analysis evaluates τe = {0.25, 0.5, 1.0, 2.0} across four Spikformer transformer blocks.
Entropy-temperature sensitivity: τe = 0.25 was adopted throughout all experiments based on the entropy-temperature sensitivity analysis.The tested temperatures were τe = 0.25, 0.5, 1.0, and 2.0.