Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Head-Tail-Aware KL Divergence in Knowledge Distillation for Spiking Neural Networks

Tianqing Zhang, Zixin Zhu, Kairong Yu, Hongwei Wang

TL;DR

The paper tackles the persistent accuracy gap between Spiking Neural Networks (SNNs) and Artificial Neural Networks (ANNs) by introducing Head-Tail-Aware KL Divergence (HTA-KL) for knowledge distillation. HTA-KL combines Forward KL and Reverse KL with adaptive head/tail weights derived from a cumulative probability mask to transfer both high- and low-probability information from a teacher ANN to an SNN student, specifically accounting for SNNs' spatio-temporal dynamics. Extensive experiments on CIFAR-10, CIFAR-100, and Tiny ImageNet show HTA-KL achieves higher accuracy with fewer timesteps, balances spike firing rates, and reduces energy consumption compared to existing KD methods. The results underscore HTA-KL's potential for efficient, high-performance neuromorphic learning and provide a practical approach to deploying SNNs on energy-constrained hardware.

Abstract

Spiking Neural Networks (SNNs) have emerged as a promising approach for energy-efficient and biologically plausible computation. However, due to limitations in existing training methods and inherent model constraints, SNNs often exhibit a performance gap when compared to Artificial Neural Networks (ANNs). Knowledge distillation (KD) has been explored as a technique to transfer knowledge from ANN teacher models to SNN student models to mitigate this gap. Traditional KD methods typically use Kullback-Leibler (KL) divergence to align output distributions. However, conventional KL-based approaches fail to fully exploit the unique characteristics of SNNs, as they tend to overemphasize high-probability predictions while neglecting low-probability ones, leading to suboptimal generalization. To address this, we propose Head-Tail Aware Kullback-Leibler (HTA-KL) divergence, a novel KD method for SNNs. HTA-KL introduces a cumulative probability-based mask to dynamically distinguish between high- and low-probability regions. It assigns adaptive weights to ensure balanced knowledge transfer, enhancing the overall performance. By integrating forward KL (FKL) and reverse KL (RKL) divergence, our method effectively align both head and tail regions of the distribution. We evaluate our methods on CIFAR-10, CIFAR-100 and Tiny ImageNet datasets. Our method outperforms existing methods on most datasets with fewer timesteps.

Head-Tail-Aware KL Divergence in Knowledge Distillation for Spiking Neural Networks

TL;DR

The paper tackles the persistent accuracy gap between Spiking Neural Networks (SNNs) and Artificial Neural Networks (ANNs) by introducing Head-Tail-Aware KL Divergence (HTA-KL) for knowledge distillation. HTA-KL combines Forward KL and Reverse KL with adaptive head/tail weights derived from a cumulative probability mask to transfer both high- and low-probability information from a teacher ANN to an SNN student, specifically accounting for SNNs' spatio-temporal dynamics. Extensive experiments on CIFAR-10, CIFAR-100, and Tiny ImageNet show HTA-KL achieves higher accuracy with fewer timesteps, balances spike firing rates, and reduces energy consumption compared to existing KD methods. The results underscore HTA-KL's potential for efficient, high-performance neuromorphic learning and provide a practical approach to deploying SNNs on energy-constrained hardware.

Abstract

Spiking Neural Networks (SNNs) have emerged as a promising approach for energy-efficient and biologically plausible computation. However, due to limitations in existing training methods and inherent model constraints, SNNs often exhibit a performance gap when compared to Artificial Neural Networks (ANNs). Knowledge distillation (KD) has been explored as a technique to transfer knowledge from ANN teacher models to SNN student models to mitigate this gap. Traditional KD methods typically use Kullback-Leibler (KL) divergence to align output distributions. However, conventional KL-based approaches fail to fully exploit the unique characteristics of SNNs, as they tend to overemphasize high-probability predictions while neglecting low-probability ones, leading to suboptimal generalization. To address this, we propose Head-Tail Aware Kullback-Leibler (HTA-KL) divergence, a novel KD method for SNNs. HTA-KL introduces a cumulative probability-based mask to dynamically distinguish between high- and low-probability regions. It assigns adaptive weights to ensure balanced knowledge transfer, enhancing the overall performance. By integrating forward KL (FKL) and reverse KL (RKL) divergence, our method effectively align both head and tail regions of the distribution. We evaluate our methods on CIFAR-10, CIFAR-100 and Tiny ImageNet datasets. Our method outperforms existing methods on most datasets with fewer timesteps.
Paper Structure (22 sections, 15 equations, 4 figures, 3 tables)

This paper contains 22 sections, 15 equations, 4 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Related Works
  3. SNN Training Methods
  4. Knowledge Distillation for SNNs
  5. Methodology
  6. Preliminaries
  7. Leaky Integrate-and-Fire Model
  8. Knowledge Distillation for SNNs
  9. FKL and RKL
  10. HTA-KL Divergence for SNNs
  11. Experiments
  12. Experimental Settings
  13. Datasets.
  14. Implementation Details.
  15. Comparison with Knowledge distillation Methods
  16. ...and 7 more sections

Figures (4)

  • Figure 1: The framework of the proposed HTA-KL for SNN training.
  • Figure 2: Impact of the Head-Tail Ratio on Accuracy. The graph shows the accuracy of ResNet-20 on CIFAR-100 at timestep 4, with different ratios between the head and tail losses. The accuracy improves as the ratio of the head loss increases, reaching optimal performance at a balanced head-tail ratio.
  • Figure 3: Spike firing rate comparison for different distillation methods on CIFAR-100. The figure shows the firing rate across different models (ResNet19, ResNet20, and VGG16) and methods (KDSNN, BKDSNN, LASNN, and HTA-KL).
  • Figure 4: t-SNE Visualization of features learned by teacher ANN and different distillation methods.