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

Kiran Nair; Rodrigue Rizk; KC Santosh

cs.LGAug 13, 2026

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

Authors: Kiran Nair, Rodrigue Rizk, KC Santosh

Organizations: USD Artificial Intelligence Research Lab Department of Computer Science University of South Dakota Vermillion, SD 57069, USA

Abstract

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.

Figures & tables

Figure 1: Overview of the proposed SAGE ( S urrogate-gradient A daptation via attention- G uided E ntropy) framework. Top: Standard Spikformer architecture consisting of a spiking patch splitting module, SSA, Multilayer Perceptron(MLP) blocks, and a classification head. Bottom: Detailed view of the SSA module. Multi-head attention (MHA) maps are generated from the spike-based query ( $Q$ ), key ( $K$ ), and value ( $V$ ) representations. Per-head attention entropy is computed from each attention map, and the standard deviation across attention heads is used as an uncertainty signal to guide adaptive surrogate-gradient modulation during training. The forward inference path remains unchanged, while the uncertainty-guided controller is active only during backpropagation.

Method	Architecture	Params (M)	Timesteps	CIFAR-10 (%)	CIFAR-100 (%)
Hybrid Training Rathi et al. (2020)	VGG-11	9.27	125	92.22	67.87
Diet-SNN Rathi and Roy (2020)	ResNet-20	0.27	10/5	92.54	64.07
STBP-tdBN Zheng et al. (2021)	ResNet-19	12.63	4	92.92	70.86
TET Deng et al. (2022)	ResNet-19	12.63	4	94.44	74.47
Spikformer Zhou et al. (2022)	Spikformer-4-384	9.32	4	95.51	78.21
RMP-SNN Han et al. (2020)	VGG-16	138.4	2048	93.63	70.93

Table 1: Comparison of state-of-the-art SNN and Spiking Transformer methods on CIFAR-10 and CIFAR-100.

Panel of Figure 2 — Figure 2: Top-1 accuracy of fixed, learnable, and SAGE surrogate gradients across datasets and simulation timesteps.

Panel of Figure 3 — Figure 3: Per-head attention entropy variability across the four SSA blocks during training. Each heatmap visualizes the standard deviation of the normalized attention entropy for each attention head over training epochs, computed from temperature-scaled attention maps ( $T=0.25$ ). Rows correspond to attention heads, columns represent training epochs.

Method	Top-1 (%)
Hybrid Training Rathi et al. (2020)	61.48
STBP-tdBN Zheng et al. (2021)	63.72
TET Deng et al. (2022)	63.72
QCFS Bu et al. (2023)	53.54
Spikformer Zhou et al. (2022)	70.24
SAGE (Ours)	60.60

Table 2: Comparison on the ImageNet-200 benchmark.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 4 — Figure 4: Evolution of Gini impurity across the four Spikformer transformer blocks during training under different attention-temperature values ( $T_{\mathrm{ent}}=0.25$ , $0.5$ , $1.0$ , and $2.0$ ). Across all blocks, Gini impurity remains highly saturated with only minor temporal fluctuations, indicating limited sensitivity to changes in attention uncertainty.

Panel of Figure 5 — Figure 5: Sensitivity analysis of the entropy temperature. Evolution of the standard deviation of normalized attention entropy across the four Spikformer blocks during training under different entropy temperatures ( $T_{\mathrm{ent}}=0.25$ , $0.5$ , $1.0$ , and $2.0$ ). Based on this analysis, $T_{\mathrm{ent}}=0.25$ was adopted throughout all experiments.

\hat{y}=f_{\theta}(x)

Appendix

Algorithm 1 SAGE Training for Spikformer

Spiking Neural Networks

Spiking neural networks process information through discrete, time-dependent neural events, aiming for efficient computation and biologically inspired learning.

See all 351 papers

A2SG:Adaptive and Asymmetric Surrogate Gradients for Training Deep Spiking Neural Networks
May 29, 2026Yechan Kang, Yongjin Kweon, Mingyeong Seo +8Neural Network Optimization Spiking Neural Networks
Center for Semiconductor Technology, Korea Institute of Science and Technology (KIST), Seoul, Republic of Korea · Department of Computer Science and Engineering, Korea University, Seoul, Republic of Korea · LIG Defense&Aerospace, Seongnam, Republic of Korea +1
Sharpness Aware Surrogate Training for Spiking Neural Networks
Mar 14, 2026Maximilian NicholsonSpiking Neural Networks Sharpness-Aware Minimization
University of Bath, United Kingdom
Advancing Direct Training for Spiking Neural Networks with Circulate-Firing Neurons and Learnable Gradients
May 14, 2026Feifan Zhou, Xiang Wei, Yang Liu +1Spiking Neural Networks Neural Network Training Dynamics
School of Artificial and Intelligence, Tianjin University, Tianjin 300072, China

May 29, 2026J/K move · Enter open · S save

A2SG:Adaptive and Asymmetric Surrogate Gradients for Training Deep Spiking Neural Networks

Yechan Kang, Yongjin Kweon, Mingyeong Seo, Sohee Park+7

Center for Semiconductor Technology, Korea Institute of Science and Technology (KIST), Seoul, Republic of Korea · Department of Computer Science and Engineering, Korea University, Seoul, Republic of Korea · LIG Defense&Aerospace, Seongnam, Republic of Korea+1

Neural Network Optimization Spiking Neural Networks Neural Network Training Dynamics

Open paper AskPDF arXiv

Training deep spiking neural networks (SNNs) remains challenging due to sharp loss landscapes and temporal inconsistency caused by surrogate gradients. To address these challenges, we propose a unified framework: adaptive and asymmetric surrogate gradients A2SG. The adaptive gradients adjust an effective window for spatio-temporal adaptation, reducing spatial gradient variation and maintaining directional consistency of gradients over time. The asymmetric gradients reflect neuronal dynamics by assigning larger gradients to neurons with higher membrane potentials, and we prove that they yield lower variation than symmetric surrogates. Our analysis further establishes a direct connection between local gradient variation and the curvature of the loss landscape, providing a principled explanation for how A2SG promotes convergence to flatter minima and improves generalization. We conduct extensive experiments on diverse models, including CNN-based and Transformer-based SNNs, across various tasks such as image classification using both static and neuromorphic datasets, as well as segmentation. The results demonstrate that A2SG consistently improves accuracy and energy efficiency, establishing it as a general and reliable solution for training deep SNNs. Our code is available at https://github.com/KIST-NCL/A2SG.git.

Abstract

Spiking Neural Networks

Explore similar work

A2SG:Adaptive and Asymmetric Surrogate Gradients for Training Deep Spiking Neural Networks

Sharpness Aware Surrogate Training for Spiking Neural Networks

Advancing Direct Training for Spiking Neural Networks with Circulate-Firing Neurons and Learnable Gradients