《Pattern Recognition》:PRF: Parallel Resonate and Fire Neuron for long sequence learning in Spiking Neural Networks
编辑推荐:
•一种解耦重置方法,将漏积分-发放(LIF)神经元的训练复杂度从O(L2)降低到O(L log L),在长序列上实现最高16.5倍的训练加速。•一种并行共振-发放(PRF)神经元,具有复数域中的振荡膜电位和可微分重置机制,能有效捕获长程依赖关系同时支持并行训练。•一种脉冲驱动
- •
一种解耦重置方法,将漏积分-发放(LIF)神经元的训练复杂度从O(L2)降低到O(L log L),在长序列上实现最高16.5倍的训练加速。
- •
一种并行共振-发放(PRF)神经元,具有复数域中的振荡膜电位和可微分重置机制,能有效捕获长程依赖关系同时支持并行训练。
- •
一种脉冲驱动的Token与通道混合器(SD-TCM)架构,集成PRF神经元,与结构化状态空间模型(SSMs)相比实现98.51%的能量降低,在长程竞技场(Long Range Arena)基准上仅下降3.19%的准确率。
引言
长序列建模是机器学习中的一个基本挑战,对大语言模型、强化学习和生成建模等应用至关重要。Transformer [1] 使用自注意力机制(具有 复杂度)进行Token混合,使用前馈网络(FFN)进行通道混合。然而,当序列长度 增加时,它们的可扩展性较差。结构化状态空间模型(SSMs)[2] 通过采用具有 推理复杂度的递归公式来解决这一问题,通过正交多项式投影 [3] 压缩输入历史。在大规模序列建模中,Mamba [4] 作为一种显著的SSMs变体,在大尺度序列建模任务中优于Transformer。然而,这两种方法仍然依赖稠密的矩阵-向量乘法,计算和能耗成本高,在资源受限环境中尤为突出。
脉冲神经网络(SNNs)[5] 通过将耗能的浮点(FP)乘累加(MAC)操作替换为稀疏的、事件驱动的累加(AC)操作 [6]、[7] 来实现节能计算。为将这一优势扩展到长序列,近期的架构层面方法将脉冲神经元集成到注意力机制和SSMs层中,产生了脉冲Transformer [8]、[9] 和脉冲SSMs [10]。虽然脉冲SSMs通过将神经元递归与SSMs的内部动力学对齐来规避Transformer的二次推理复杂度,但它们仍然依赖SSMs层内的稠密矩阵-向量乘法 [11]。这损害了脉冲的能效优势。此外,建模长程依赖的能力仍然依赖于注意力或SSMs架构 [12],而非神经元本身。因此,脉冲神经元的固有递归和轻量级计算优势仍未得到充分利用。有必要重新审视神经元层面的机制,充分利用脉冲神经元内在支持长程依赖建模的动力学特性,同时具有最小的能耗开销。
使用脉冲神经元实现长序列学习引入了若干基本挑战。(i) 通过时间反向传播(BPTT)是SNNs常用的训练方法 [13]、[14]。然而,其训练时间随序列长度呈二次增长 [15]。尽管存在高效的SNN训练方法 [16],但在长序列上它们的性能仍不如BPTT [17]。(ii) 常用的脉冲神经元,如漏积分-发放(LIF)模型,难以捕获长程依赖关系。先前工作通过更复杂的动力学增强LIF神经元 [18]、[19],但往往以牺牲训练效率和在序列长度及模型规模上的较差可扩展性为代价。(iii) 为提高训练速度,近期研究探索并行化方法。然而,它们往往牺牲了推理效率。例如,PSN系列 [20] 需要缓存多个输入步骤,PMSN [21] 仍然依赖稠密矩阵操作,这些都阻碍了在神经形态硬件上的部署。因此,实现一种支持长序列建模(例如 k)、支持并行训练、并保持高效推理(即无需缓存,仅逐元素递归)的脉冲神经元仍然是一个开放且尚未充分探索的挑战。
在本工作中,我们提出了一种神经元层面的长序列建模解决方案,平衡了训练可扩展性和推理效率,如图1所示。我们的贡献如下:
(1) 为加速训练过程,我们提出了一种解耦重置方法,允许并行训练同时保持与序列模型等效的脉冲行为,实现了211.82 的反向传播加速。该方法揭示了一个见解:软重置在数学上等价于一种自适应阈值机制。
(2) 为有效提取长程依赖关系,我们提出了具有复数域中振荡膜电位的并行共振-发放(PRF)神经元,利用自适应和可微分的重置机制。
(3) 为将PRF扩展到更复杂的长程学习任务,我们进一步设计了一个轻量级的脉冲驱动时间-通道混合器(SD-TCM)模块。SD-TCM作为评估神经元长序列建模能力的最小架构。集成到SD-TCM中,PRF相比S4实现了98.51%的能量降低,在长程竞技场基准上仅下降3.19%的准确率。
各章节摘要
推理高效的神经元
神经形态硬件偏好轻量级操作,如逐元素乘法或累加,而非耗能的矩阵乘法和激活函数 [22]。众多脉冲神经元模型已被开发以改善时间动力学同时遵循这些约束。例如,PLIF [23] 和GLIF [24] 引入参数可塑性;TC-LIF [25] 及其改进版TS-LIF [26] 模拟双室交互,LM-HT [27] 进一步利用可学习的多层级
问题形式化
为实现高效的并行训练和长序列建模,同时严格保持推理效率,我们重新审视了广泛使用的LIF神经元。我们首先回顾LIF模型并重新思考并行训练和长序列学习的两个关键挑战。
LIF神经元是一种广泛使用的脉冲模型 [13]。它模拟生物神经元,膜电位 的动态遵循式(1),其中 是输入电流,、 和 是膜
方法
在本节中,我们提出我们的方法来应对LIF中的两个关键挑战:实现并行训练和改善长程依赖关系建模。为支持并行训练,我们将重置机制与积分过程解耦,在保持与原始序列输出等效的同时实现高效计算。为进一步增强长程建模,我们提出了并行共振-发放神经元,将重置过程作为虚部纳入
实验
为评估PRF的效率和有效性,我们在各种脉冲神经元模型上评估训练/推理效率和长序列建模能力。然后我们分析解耦重置方法带来的训练加速。进一步地,我们考察PRF神经元捕获长程依赖的能力,证明核振荡增强了性能和梯度稳定性。最后,PRF实现了与S4相当的性能,同时将能耗降低了超过98.51%
讨论
除通用长序列基准外,高效的长程依赖建模在计算密集型模式识别任务中越来越重要。例如,高分辨率高光谱图像超分辨率需要建模长程光谱-空间依赖关系 [43];细粒度视觉识别最近采用了Mamba风格的扫描来捕获判别性长程结构 [44];以及空间-时间建模任务,如矩阵分解和
结论
本研究解决了脉冲神经网络中的并行化和长序列挑战。我们的解耦重置方法降低了训练复杂度,实现了最高16.50 的加速。我们通过振荡膜动力学构建了PRF神经元以捕获长程依赖。集成到脉冲驱动的架构中,PRF在长程竞技场基准上实现了98.51%的理论能量降低,准确率仅下降3.19%。总体而言,我们的方法提供了一种可扩展的解决方案,平衡了
CRediT作者贡献声明
黄玉龙(Yulong Huang):写作 – 原稿撰写,软件,方法论,概念化。刘遵昌(Zunchang Liu):写作 – 审阅与编辑,资源。林晓鹏(Xiaopeng Lin):写作 – 审阅与编辑。任宏伟(Hongwei Ren):写作 – 审阅与编辑。黄宏向(Hongxiang Huang):写作 – 审阅与编辑。傅浩天(Haotian Fu):写作 – 审阅与编辑。周悦(Yue Zhou):写作 – 审阅与编辑。邢宏(Hong Xing):写作 – 审阅与编辑。谢泽科(Zeke Xie):写作 – 审阅与编辑,指导。程步军(Bojun Cheng):写作 – 审阅与编辑,指导,资金支持
利益冲突声明
作者声明他们没有任何已知的可能影响本文所报告工作的竞争性经济利益或个人关系。
致谢
本工作部分由国家自然科学基金青年科学基金(批准号62305278)、广东省基础与应用基础研究基金(编号2025A1515011758)和"1+1+1":香港科技大学-香港科技大学(广州)跨校区科研合作"1+1+1"联合资助计划(G042)支持。
黄玉龙|刘遵昌|林晓鹏|任宏伟|黄宏向|傅浩天|周悦|邢宏|谢泽科|程步军