PaperlayerView exact revision on arXiv

Source-linked AI summary

Temporal Efficient Training of Spiking Neural Network via Gradient Re-weighting

Shikuang Deng, Yuhang Li, Shanghang Zhang, Shi Gu

arXiv:2202.11946v3cs.NEcs.AI

TL;DR

Deep SNNs are difficult to train because non-differentiable spikes and mismatched surrogate gradients can produce poor generalization. The paper introduces TET, which optimizes each time point to recover flatter-minimum behavior and improve temporal scalability. TET consistently outperforms prior methods across reported datasets, including 83.17% top-1 accuracy on DVS-CIFAR10 and nearly halved training time through TIT.

  • Problem

    Non-differentiable spikes hinder gradient-based training of deep SNNs, while surrogate-gradient direct training can mismatch the SNN loss landscape and yield poor generalizability.

  • Method

    TET optimizes every moment’s pre-synaptic inputs to compensate for surrogate-gradient momentum loss and find flatter minima, while TIT initializes SNNs with shorter simulations.

  • Results

    TET consistently outperforms SDT across static and neuromorphic datasets, including 83.17% top-1 accuracy on DVS-CIFAR10, over 10% above the current state of the art.

  • Takeaways & Limitations

    TET improves SNN temporal scalability and enables TIT to reduce training-time consumption by almost half.

Abstract

from arXiv · show

Recently, brain-inspired spiking neuron networks (SNNs) have attracted widespread research interest because of their event-driven and energy-efficient characteristics. Still, it is difficult to efficiently train deep SNNs due to the non-differentiability of its activation function, which disables the typically used gradient descent approaches for traditional artificial neural networks (ANNs). Although the adoption of surrogate gradient (SG) formally allows for the back-propagation of losses, the discrete spiking mechanism actually differentiates the loss landscape of SNNs from that of ANNs, failing the surrogate gradient methods to achieve comparable accuracy as for ANNs. In this paper, we first analyze why the current direct training approach with surrogate gradient results in SNNs with poor generalizability. Then we introduce the temporal efficient training (TET) approach to compensate for the loss of momentum in the gradient descent with SG so that the training process can converge into flatter minima with better generalizability. Meanwhile, we demonstrate that TET improves the temporal scalability of SNN and induces a temporal inheritable training for acceleration. Our method consistently outperforms the SOTA on all reported mainstream datasets, including CIFAR-10/100 and ImageNet. Remarkably on DVS-CIFAR10, we obtained 83$\%$ top-1 accuracy, over 10$\%$ improvement compared to existing state of the art. Codes are available at \url{https://github.com/Gus-Lab/temporal_efficient_training}.

1 INTRODUCTION

SNNs offer energy-efficient, fast inference but remain difficult to train directly because spikes are non-differentiable and surrogate-gradient optimization can generalize poorly. The paper introduces TET to optimize temporal outputs and reports improvements across static and neuromorphic datasets.

  • SNNs avoid multiplication during inference through binary spike transmission, enabling energy-saving and fast computation on neuromorphic hardware.
  • Surrogate-gradient direct training can mismatch the SNN loss landscape and become trapped in local minima with low generalizability.
  • Deep SNNs trained from scratch lag behind ANNs, while training memory and time consumption increase linearly with simulation time.
  • TET optimizes every moment’s pre-synaptic inputs instead of directly optimizing integrated potential, improving time scalability and enabling TIT initialization with shorter simulations.
  • 83.17% top-1 accuracy on DVS-CIFAR10 exceeded the current state-of-the-art result by over 10%.

2 RELATED WORK

SNN training research primarily uses ANN-to-SNN conversion or direct training from scratch. Conversion can achieve high accuracy but incurs latency and is unsuitable for neuromorphic data, whereas direct training uses surrogate gradients with fewer time steps but higher training costs.

  • SNN training commonly follows either ANN-to-SNN conversion or direct training from scratch.
  • ANN-to-SNN Conversion: ANN-to-SNN conversion transfers a high-performing ANN into an SNN and adjusts parameters layer by layer to match ANN activations.
  • ANN-to-SNN Conversion: Converted methods can reduce inference latency with specialized normalization and calibration techniques, but they significantly extend latency and are unsuitable for neuromorphic data.
  • Direct training: Direct training treats SNNs as recurrent networks, uses backpropagation through time with surrogate gradients, and achieves high accuracy with few time steps but higher training costs.

3 PRELIMINARY

The paper models neurons with leaky integrate-and-fire dynamics, uses surrogate gradients for spatial-temporal backpropagation, and applies temporal-spatial batch normalization. Classification uses average membrane potential rather than final-layer spikes.

  • 3.1 ITERATIVE LIF MODEL: The LIF model updates membrane potential iteratively using a leaky factor and pre-synaptic input, which equals synaptic weight multiplied by the spiking input.
  • 3.1 ITERATIVE LIF MODEL: When membrane potential exceeds threshold Vth, the neuron fires a spike and resets its membrane potential to zero.
  • 3.1 ITERATIVE LIF MODEL: The experiments use starting membrane u(0)=0, threshold Vth=1, and leaky factor τ=0.5.
  • 3.1 ITERATIVE LIF MODEL: Average membrane potential serves as the classification index, with cross-entropy loss used for training instead of last-layer spike frequency.
  • 3.2 SURROGATE GRADIENT: Spatial-temporal backpropagation treats the SNN as an RNN, while surrogate gradients replace the non-differentiable step-function derivative.
  • 3.2 SURROGATE GRADIENT: The method uses triangular surrogate gradients, whose constraint factor γ determines the sample range where gradients activate.
  • 3.3 BATCH NORMALIZATION: Threshold-dependent batch normalization normalizes pre-synaptic inputs across spatial and temporal dimensions.

4 METHODOLOGY

TET replaces standard direct training’s whole-output optimization with moment-wise output constraints and gradient re-weighting, targeting flatter minima and improved temporal scalability. Its training pipeline combines the TET loss with an MSE regularizer, while inference rules remain unchanged.

  • 4.1 FORMULA OF TRAINING SNN WITH SURROGATE GRADIENTS: Standard direct training optimizes the network’s integrated output, whereas TET constrains each moment’s pre-synaptic output toward the target distribution.TET introduces a temporal loss that applies supervision at every simulation time point.
  • 4.2 CONVERGENCE OF GRADIENT DESCENT FOR SDT V.S. TET: Surrogate-gradient mismatch can dissipate accumulated momentum near local minima, preventing SDT from searching for flatter minima with better generalizability.The mismatch arises because the spiking activation is a step function while the surrogate gradient is bounded by integral constraints.
  • 4.2 CONVERGENCE OF GRADIENT DESCENT FOR SDT V.S. TET: TET reweights the contribution of each output derivative, increasing gradient norms near sharp minima and driving optimization toward flatter minima.Its convergence condition relies on the per-time-point derivative becoming small rather than the error term vanishing at every moment.
  • 4.2 CONVERGENCE OF GRADIENT DESCENT FOR SDT V.S. TET: LSDT is upper bounded by LTET, so driving LTET close to zero also makes the original standard-training loss approach zero.The bound follows from the arithmetic mean-geometric mean and Jensen inequalities.
  • 4.3 TIME INHERITANCE TRAINING: An MSE regularizer confines each moment’s output to reduce the risk that an outlier time point dominates the integrated SNN output.The regularization target is φ, set to Vth in the experiments, and λ controls its contribution to the total loss.
  • 4.3 TIME INHERITANCE TRAINING: TET’s moment-wise optimization enables temporal extension, while TIT first trains at a short simulation length and then fine-tunes at the target length.The paper reports that TIT improves accuracy over training from scratch and significantly saves training time.

5 EXPERIMENTS

Experiments validate TET against SDT, examine its loss-landscape and temporal-scaling behavior, and compare performance across static and neuromorphic datasets. TET favors flatter minima, supports temporal inheritance, improves efficiency, and reports gains across CIFAR, ImageNet, and DVS-CIFAR10.

  • Model validation and ablation study: TET increases ResNet-19 accuracy over SDT with surrogate gradients by 3.25% at simulation time 4 and 3.53% at simulation time 6.The comparison uses CIFAR100 and three training runs.
  • Model validation and ablation study: TET finds flatter local minima than SDT, while its loss is harder to minimize around SDT’s sharp local minima.The landscapes compare LSDT and LTET around local minima.
  • Model validation and ablation study: Changing from SDT to TET after epoch 200 rapidly reduces both LTET and LSDT on the test set.The experiment uses VGGSNN trained for 300 epochs on DVS-CIFAR10.
  • Model validation and ablation study: Increasing simulation length generally raises accuracy until performance slightly decreases, while TIT inheritance outperforms scratch training at T=4 and T=6.ResNet-19 on CIFAR100 is changed from initial lengths 2, 3, or 4 to lengths up to 8 without finetuning.
  • Model validation and ablation study: TET yields higher network efficiency across simulation times when performance is evaluated against energy consumption.The energy comparison uses 0.9pJ for addition and 4.6pJ for multiplication in 45nm CMOS technology.
  • Comparison to existing works: TET improves reported performance across CIFAR, ImageNet, and DVS-CIFAR10, reaching 83.17% accuracy on DVS-CIFAR10.Reported gains include more than 3% on CIFAR100, 1.07% on Spiking-ResNet-34, and 0.96% on SEW-ResNet34.

6 CONCLUSION

The paper identifies poor generalization in direct SNN training and proposes TET to address it. TET also improves temporal scalability, enabling TIT to nearly halve training time.

  • TET targets the generalization problem in direct SNN training, where training performance is strong but test performance is poor.
  • TET consistently achieves better performance than the standard direct training process across the reported experiments.
  • TET significantly improves the time-scalability robustness of SNNs.
  • TIT uses TET’s temporal scalability to reduce training-time consumption by almost a half.

A.1 DATASET AND TRAINING DETAIL

The experiments use distinct dataset, architecture, preprocessing, optimizer, and simulation-time settings for CIFAR, ImageNet, and DVS-CIFAR10.

  • CIFAR: CIFAR10 and CIFAR100 use ResNet-19, 300 initial training epochs, simulation length T = 2, Adam optimization, and standard crop-and-flip augmentation.
  • ImageNet: ImageNet experiments use more than 1250k training images, 50k validation images, 224×224 crops, standard training augmentation, and either SEW-ResNet34 or Spiking-ResNet34.
  • ImageNet: ImageNet training uses SGD with 0.9 momentum, weight decay 4e −5, learning rate 0.1, cosine decay, and T = 4 for SEW-ResNet34.
  • DVS-CIFAR10: DVS-CIFAR10 is converted into temporally accumulated blocks, split into 9k training and 1k test images, reduced to 48×48, and augmented with flips and random rolls.

A.2 LSDT LOSS LANDSCAPE OF RESNET-19

On CIFAR100, the SDT solution lies in a sharp loss-landscape region, whereas the TET solution lies in a smoother region associated with improved generalization.

  • TET finds a much smoother local-minimum region than SDT for ResNet-19 on CIFAR100.
  • The SDT local-minimum region is described as very sharp around its minimum.
  • The comparison uses classification-loss landscapes from different training approaches for ResNet-19 on CIFAR100.

A.3 EFFECT OF LMSE

The LMSE regular term improves performance across CIFAR100 and DVS-CIFAR10, but its useful strength depends on the dataset and can affect early ImageNet training.

  • LMSE improves the performance of ResNet-19 on CIFAR100 and VGGSNN on DVS-CIFAR10 across the tested λ levels.
  • Larger λ values are more suitable for CIFAR100, whereas smaller λ values are more suitable for DVS-CIFAR10.
  • A large λ can harm the early training phase on ImageNet, particularly when zero-initialization is not performed.
  • The reported λ comparison covers ResNet-19 on CIFAR100 and VGGSNN on DVS-CIFAR10.

A.4 STATISTICAL RESULTS

The statistical results link overall SNN performance to moment-wise accuracy and compare TET with standard training across temporal-scaling experiments. TET shows stronger moment-wise classification ability and lower accuracy variance when simulation length increases.

  • Overall SNN performance is highly positively associated with the average accuracy of each moment.Figure 7 compares standard training with TET using green and red dots, respectively.
  • TET gives membrane potential increments at each moment higher classification ability than SDT on ResNet19 trained on CIFAR100.The networks were trained with initial simulation lengths of 3 or 4 before expansion to 8.
  • Increasing simulation length causes SNN accuracy to first increase and then stabilize within a certain area.This pattern is reported for both SDT and TET in the temporal scalability comparison.
  • 1.80 accuracy variance for TET versus 11.13 for SDT after increasing simulation length.The comparison expands ResNet19 SNN simulation length to 64 after training with initial lengths of 2, 3, or 4.
  • TET initialization requires only a small simulation length, whereas SDT requires a sufficiently large simulation or enough finetuning epochs.This interpretation follows the reported variance comparison after extending simulation length.
Loading 2202.11946v3…