通过输出和多个中间特征上的离散知识蒸馏训练低延迟脉冲神经网络

《Neurocomputing》：Training low-latency spiking neural networks through discrete knowledge distillation over output and multiple intermediate features

　　

# 基于多层知识蒸馏的脉冲神经网络方法

## 摘要

脉冲神经网络（SNNs）作为一种有前景的节能替代方案，正逐步兴起，其利用事件驱动机制可实现高效的人工脑芯片部署。然而，如何在保证精度不受损的前提下实现低延迟推理仍然具有挑战性，因为SNNs通常需要较多的时间步才能进行有效处理。为此，我们提出了一种多层SNN知识蒸馏方法，同时作用于输出层和多个中间特征，从而实现低延迟的SNN训练。据我们所知，这是首项将同时针对输出层和多个中间特征的多层SNN知识蒸馏应用于低延迟SNN训练的工作，其核心依赖于一种专为离散脉冲表示而设计的基于概率分布的脉冲特征蒸馏方法。具体而言，我们首先训练一个高性能ANN，然后通过ANN-to-SNN转换方法获得一个在长时间步下与ANN性能匹配的SNN。随后，我们采用循环式教师-学生蒸馏过程：初始的长时间步SNN作为教师模型，引导训练短时间步的学生SNN。收敛后，该学生SNN成为新的教师模型，监督一个更短时间步的学生SNN。这一循环通过多层SNN知识蒸馏逐步压缩时间维度，直到获得一个用最少时间步即可匹配ANN性能的紧凑SNN。在这种多层SNN知识蒸馏中，教师和学生模型均为SNN，我们同时使用输出logits和中间特征进行全面的知识传递，但SNN的离散特性使得直接特征蒸馏不可行。为此，我们提出了一种基于概率分布的脉冲特征蒸馏方法，将脉冲序列转化为连续的概率分布，并最小化教师与学生之间的分布散度。在公开数据集上的实验表明，我们的方法显著减少了推理时间步，同时保持了有竞争力的精度。

## 引言

人工神经网络（ANNs）在各领域的显著成功正日益受到其高昂的计算成本和能耗的制约[11], [30]。因此，提高计算和能源效率已成为学术界和工业界的首要关注点。作为第三代神经网络，脉冲神经网络（SNNs）被视为一种有前景的低功耗替代方案，可能突破上述局限，因而获得了广泛关注[36]。尽管SNNs与ANNs在结构上相似，但它们在信息表示和计算机制上存在根本差异。不同于采用连续激活值的ANNs，SNNs利用离散脉冲事件进行神经元间通信，采用事件驱动范式实现稀疏计算。这一特性使SNNs能够实现显著低于ANNs的能耗，并特别适合部署在人工脑硬件上[9], [15]。

SNNs中脉冲序列的离散特性限制了ANNs中标准反向传播算法的直接应用。这催生了三种主要训练方法的发展：时空反向传播（STBP）[33]、脉冲时序依赖可塑性（STDP）[10]以及ANN-to-SNN转换[1]。其中，STBP和STDP在训练深层架构时常面临重大困难，且可能无法达到有竞争力的精度[3], [16]。相比之下，ANN-to-SNN转换通过先训练常规ANN再将其转换为等效SNN，在很大程度上规避了直接SNN方法带来的训练不稳定性问题，从而有利于部署更深、更复杂的网络。后续研究聚焦于优化转换技术，以充分利用这些优势并应对固有挑战。例如，Rathi等人[27]实施了一种混合训练方法，将ANN-to-SNN转换与脉冲时序依赖反向传播（STDB）相结合，使用转换后SNN的参数进行初始化，然后通过STDB微调以减少推理延迟。类似地，Zhou等人[45]设计了Spiking Transformer框架，引入了一种脉冲自注意力（SSA）机制，利用脉冲形式的查询、键和值矩阵并消除softmax操作，实现适用于SNNs的高效无乘法计算。

然而，ANN-to-SNN转换中仍存在一个根本性局限：该过程依赖于将连续激活值转换为离散脉冲序列，而实现高精度通常需要较长的时间窗口才能充分近似原始ANN的激活值。对大量时间步的需求直接与SNN部署的能源效率核心目标相矛盾，因为更长的脉冲序列会增加动态功耗并拖慢推理速度。近期的努力旨在通过提高脉冲信号在有限时间步内的表达能力来提升转换的效率和精度。通过改进数据编码方案或神经元模型，这些方法提高了脉冲序列的信息效率并减小了近似误差，从而减少了高性能推理所需的时间步数。Yang等人[37]提出了NeuBridge，将量化激活值与脉冲神经元桥接以实现ANN-SNN转换，使得在有限时间步下实现高精度推理。Li等人[22]提出了一种注入噪声的积分-发放脉冲神经元，在膜电位充电过程中注入高斯噪声，有效地实现了随机阈值机制，改善了低时间步条件下的放电动态和性能。

现有方法尚未将时间压缩推向实际极限，因此，对SNNs时间维度的系统性压缩以实现超低延迟推理仍然是一个关键且未被充分探索的挑战。为此，我们提出了一种多层SNN知识蒸馏方法，同时针对输出logits和多个中间特征进行知识传递。据我们所知，这是首项将同时针对输出层和多个中间特征的多层SNN知识蒸馏应用于低延迟SNN训练的工作，其核心依赖于一种专为离散脉冲表示而设计的基于概率分布的脉冲特征蒸馏方法。为克服脉冲的离散和事件驱动特性导致直接特征蒸馏不可行的问题，我们设计了一种基于概率分布的脉冲特征蒸馏方法。

具体而言，我们的方法首先训练一个高性能ANN，然后通过有效的ANN-to-SNN转换方法将其转换为在长时间步下运行的高性能SNN。随后，采用循环式教师-学生蒸馏过程：初始的长时间步SNN作为教师模型，引导短时间步学生SNN模型的训练。收敛后，该学生SNN模型成为新的教师模型，监督一个更短时间步的学生SNN模型。这一循环通过多层SNN知识蒸馏逐步压缩时间维度，确保稳定的精度传递并缓解激进单步压缩通常导致的精度损失。该过程持续进行，直到获得一个用最少时间步即可匹配ANN性能的紧凑SNN。在多层SNN知识蒸馏过程中，学生SNN不仅学习教师SNN输出logits中的知识，还学习其中间特征中的知识；教师和学生均实现为SNNs。此外，当从教师SNN的中间特征学习时，学生SNN同时学习多个中间特征，而非仅关注特定特征层。这使得学生SNN的输出logits和特征表示能够尽可能逼近教师SNN。然而，在多层SNN知识蒸馏过程中，SNNs中离散的脉冲活动使得直接特征蒸馏不切实际。为此，我们提出了一种基于概率分布的脉冲特征蒸馏方法，实现了从脉冲序列的有效知识传递。该方法将脉冲序列转化为连续的特征概率分布，并最小化教师与学生模型之间的分布散度，从而实现从SNNs离散特征表示中的知识蒸馏。我们的方法的关键创新在于两点：

1. 我们首次将同时针对输出logits和多个中间特征的多层SNN知识蒸馏应用于低延迟SNN训练，其中长时间步教师SNN通过基于概率分布的脉冲特征蒸馏监督短时间步学生SNN，从而得到一个即使时间步极低也能实现与ANNs相当性能的SNN。
2. 创新性地设计了一种适用于SNNs的基于概率分布的脉冲特征蒸馏方法，实现了从长时间步教师到短时间步学生的离散脉冲特征的有效知识传递，从而直接在SNNs离散特征表示之间实现知识蒸馏。

本文其余部分组织如下：第2节回顾SNN架构和训练方法的相关工作。第3节介绍我们的方法及其核心组件。第4节展示实验结果，第5节进行讨论和分析。最后，第6节总结了我们的结论。

## 相关工作

尽管ANNs在各领域取得了显著成功，但其高昂的计算成本和能耗已成为主要瓶颈，促使人们日益关注SNNs作为节能替代方案，近年来也取得了重大进展。不同于依赖连续激活值的ANNs，SNNs的根本区别在于通过离散脉冲事件进行通信并以事件驱动方式运行，从而实现稀疏高效的计算[26]。这种固有特性使得...

## 方法

本节详细阐述了所提出的用于训练低延迟SNNs的多层SNN知识蒸馏方法。该方法的结构依次为：首先介绍整体框架，然后是核心技术——基于概率分布的脉冲特征蒸馏，接着是针对输出层和多个中间特征的多层SNN知识蒸馏，最后是通过循环蒸馏实现的渐进式时间压缩。表1展示了所使用重要参数的描述...

## 实验

本节提供所提出的用于训练低延迟脉冲神经网络的多层SNN知识蒸馏方法的实验结果和分析。首先，我们详细介绍实验配置，包括神经元模型、网络架构和超参数设置。随后，我们介绍几个公开数据集，包括传统图像分类任务和人工脑事件视觉任务。然后我们呈现广泛的实验结果...

## 讨论

针对输出层和多个中间特征的多层SNN知识蒸馏方法在多种数据集和架构上均成功实现了SNNs的低延迟推理，同时保持了有竞争力的精度。除直接性能指标外，还有三个更广泛的启示值得进一步讨论：与竞争性的低延迟SNN方法相比的训练时间成本、推理时激进时间压缩带来的能效收益，以及方法的通用性...

## 结论

本文提出了一种新颖的多层SNN知识蒸馏方法，用于训练低延迟脉冲神经网络。据我们所知，这是首项将同时针对输出层和多个中间特征的多层SNN知识蒸馏应用于低延迟SNN训练的工作，其核心依赖于一种专为离散脉冲表示而设计的基于概率分布的脉冲特征蒸馏方法。通过同时利用输出logits和多个中间特征进行知识传递...

## CRediT作者贡献声明

葛步锋：撰写—初稿、可视化、验证、软件、方法论、调查研究、正式分析。杜岚：撰写—审阅与编辑、监督、资源、项目管理、方法论、调查研究、概念设计。郭宇晨：撰写—审阅与编辑、方法论、调查研究。宋佳伦：方法论。杨兴瑞：可视化。刘欢：方法论。

## 利益冲突声明

作者声明不存在可能影响本文所报告工作的已知竞争性经济利益或个人关系。

## 致谢

本研究部分由国家自然科学基金资助，项目编号为U24B20137、U21B2039和62201431。

葛步锋于2024年获得中国西安电子科技大学电子信息工程学士学位。目前他在西安电子科技大学雷达信号处理国家重点实验室攻读信号处理博士学位。他的研究兴趣包括雷达自动目标识别和脉冲神经网络。

葛步锋 | 杜岚 | 郭宇晨 | 宋佳伦 | 杨兴瑞 | 刘欢