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

📄 arXiv: 2504.20445v2 📥 PDF

作者: Tianqing Zhang, Zixin Zhu, Kairong Yu, Hongwei Wang

分类: cs.AI

发布日期: 2025-04-29 (更新: 2025-05-16)

备注: Accepted by IJCNN2025

💡 一句话要点

提出Head-Tail-Aware KL散度，用于SNN知识蒸馏，提升模型泛化能力。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 脉冲神经网络 知识蒸馏 KL散度 Head-Tail Aware 模型压缩

📋 核心要点

  1. 现有基于KL散度的知识蒸馏方法在SNN中存在不足，过度关注高概率区域，忽略低概率区域，导致泛化能力受限。
  2. 提出Head-Tail Aware KL (HTA-KL)散度，通过累积概率掩码动态区分高低概率区域，并自适应地调整权重，平衡知识转移。
  3. 实验结果表明，在CIFAR-10、CIFAR-100和Tiny ImageNet数据集上，该方法在更少的时间步长下优于现有方法。

📝 摘要（中文）

脉冲神经网络(SNNs)作为一种节能且生物可解释的计算方法备受关注。然而，由于现有训练方法的局限性和固有的模型约束，与人工神经网络(ANNs)相比，SNNs通常存在性能差距。知识蒸馏(KD)已被探索为一种将知识从ANN教师模型迁移到SNN学生模型以弥合这一差距的技术。传统的KD方法通常使用Kullback-Leibler (KL)散度来对齐输出分布。然而，传统的基于KL的方法未能充分利用SNN的独特特性，因为它们倾向于过度强调高概率预测而忽略低概率预测，导致次优的泛化。为了解决这个问题，我们提出了一种新的SNNs的KD方法，即Head-Tail Aware Kullback-Leibler (HTA-KL)散度。HTA-KL引入了一个基于累积概率的掩码，以动态区分高概率区域和低概率区域。它分配自适应权重以确保平衡的知识转移，从而提高整体性能。通过整合前向KL (FKL)和反向KL (RKL)散度，我们的方法有效地对齐了分布的头部和尾部区域。我们在CIFAR-10、CIFAR-100和Tiny ImageNet数据集上评估了我们的方法。在大多数数据集上，我们的方法以更少的时间步长优于现有方法。

🔬 方法详解

问题定义：现有基于KL散度的知识蒸馏方法在训练SNN时，容易过度关注教师网络输出分布中的高概率区域（head），而忽略低概率区域（tail）。这种不平衡的知识迁移导致学生网络（SNN）的泛化能力下降，无法充分学习教师网络（ANN）的知识。

核心思路：论文的核心思路是设计一种新的KL散度计算方法，即Head-Tail Aware KL (HTA-KL)散度，能够自适应地调整高低概率区域的权重，从而平衡知识迁移。通过引入累积概率掩码，动态区分head和tail区域，并分别进行优化，使得SNN能够更好地学习ANN的知识，提升泛化能力。

技术框架：HTA-KL知识蒸馏框架主要包含以下几个阶段：1) 使用ANN作为教师网络进行预训练；2) 使用SNN作为学生网络，并初始化参数；3) 计算ANN的输出分布；4) 使用HTA-KL散度作为损失函数，指导SNN的学习，其中HTA-KL散度结合了前向KL (FKL)和反向KL (RKL)散度，分别用于对齐head和tail区域；5) 通过反向传播更新SNN的参数。

关键创新：该方法最重要的创新点在于提出了Head-Tail Aware KL (HTA-KL)散度。与传统的KL散度相比，HTA-KL能够动态地关注输出分布的不同区域，并自适应地调整权重，从而平衡知识迁移。此外，结合使用FKL和RKL散度，能够更全面地对齐教师网络和学生网络的输出分布。

关键设计：HTA-KL散度的关键设计在于累积概率掩码的计算和自适应权重的分配。累积概率掩码用于区分head和tail区域，通常设置一个阈值来划分。自适应权重的分配则根据head和tail区域的重要性进行调整，例如，可以对tail区域赋予更高的权重，以鼓励SNN学习低概率事件的知识。此外，FKL和RKL散度的权重也需要进行调整，以平衡两个方向的知识迁移。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在CIFAR-10、CIFAR-100和Tiny ImageNet数据集上，提出的HTA-KL方法在大多数情况下优于现有的知识蒸馏方法。例如，在CIFAR-10数据集上，使用更少的时间步长，HTA-KL方法能够达到更高的分类精度。这表明HTA-KL方法能够更有效地将知识从ANN教师网络迁移到SNN学生网络，提升SNN的性能。

🎯 应用场景

该研究成果可应用于各种需要节能高效计算的场景，例如边缘计算、移动设备和物联网设备。通过将知识从大型ANN模型迁移到小型SNN模型，可以在资源受限的设备上部署高性能的神经网络，实现更智能化的应用，例如智能监控、自动驾驶和语音识别等。

📄 摘要（原文）

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.