近日,一篇名为《SpikingMamba: Towards Energy-Efficient Large Language Models via Knowledge Distillation from Mamba》的学术论文发表于Transactions on Machine Learning Research(发表时间:01/2026)。该研究由来自香港科技大学(广州)、香港城市大学、南方科技大学、华东师范大学和华为技术有限公司等机构的研究人员共同完成,其中Yulong Huang、Jianxiong Tang和Chao Wang为并列第一作者,Bojun Cheng和Luziwei Leng为通讯作者。这篇论文在大型语言模型(Large Language Models, LLMs)的效率优化领域提出了一项创新性工作,旨在解决Transformer架构及新兴的Mamba架构在能效方面面临的挑战,通过引入脉冲神经网络(Spiking Neural Networks, SNNs)技术,构建了首个基于知识蒸馏的、具备线性时间复杂度的十亿参数级别脉冲语言模型。
一、学术背景与目标
大型语言模型在诸多任务上展现了卓越性能,但其巨大的计算与能耗限制了在边缘设备等资源受限环境下的部署。传统的Transformer架构因其自注意力机制的二次方复杂度而效率受限。近年来提出的Mamba架构,采用选择性状态空间模型(Selective State Space Model, SSM)替代注意力机制,实现了线性时间复杂度的序列建模,并消除了推理时的键值缓存开销,为长上下文和边缘场景带来了效率提升。然而,Mamba模型仍严重依赖密集矩阵乘法,其能耗问题在电池供电的嵌入式设备上依然突出。
与此同时,脉冲神经网络作为一种生物启发式计算模型,因其事件驱动、稀疏激活的特性,在理论上能够将密集的乘积累加运算转换为稀疏的累加运算,从而显著提升能效。然而,现有的脉冲语言模型(SNN-based LLMs)往往面临两难选择:要么为了效率而牺牲性能,要么为了恢复性能需要耗费高昂成本从头训练大模型。
基于以上背景,本研究提出了一项明确的目标:开发一种名为SpikingMamba的高能效脉冲语言模型。该模型旨在利用Mamba的架构效率,并结合SNN的节能优势,通过知识蒸馏等高效训练策略,在最大限度保留Mamba模型零样本推理能力的同时,显著降低模型的推理能耗。研究的核心科学问题是:如何设计有效的脉冲神经元和训练框架,将预训练大语言模型(尤其是Mamba)的能力高效迁移到脉冲架构上,以极小的性能损失换取显著的能效收益。
二、详细研究方法与流程
本研究包含一个完整的研究流程,主要分为模型设计、训练框架构建和实验验证三个阶段,研究对象为不同参数规模(130M和1.3B)的语言模型。
第一阶段:SpikingMamba模型架构设计 研究首先分析了Mamba2模型的构成,发现超过90%的参数位于输入和输出的线性投影层,这些层是计算密集型的核心部分。因此,SpikingMamba的核心思想是在这些投影层之前插入脉冲神经元,将密集的矩阵向量乘法转换为由稀疏脉冲触发的行累积操作。这一设计的关键在于两个创新组件:
有符号整数泄漏积分发放(Signed Integer Leaky Integrate-and-Fire, SI-LIF)神经元:为了解决传统LIF神经元无法保留语义极性(负激活)以及现有整数LIF(I-LIF)神经元丢失幅度信息的问题,研究者设计了SI-LIF神经元。该神经元在训练时输出有符号的整数值(范围在[-d, d]),既能保留激活的符号(正负)又能保留相对幅度信息,形成了一个三值脉冲方案。在推理时,通过内部的微时间步动力学,SI-LIF可以完全等价地重构出训练时的整数输出,并转化为事件驱动的二进制脉冲序列进行计算,确保了稀疏计算优势。
平滑梯度补偿路径:为了缓解脉冲量化带来的表征保真度损失,研究引入了一个仅在训练时使用的辅助路径。该路径使用一个可微的模仿函数(如d × tanh(x))来模拟SI-LIF神经元输出的动态范围,并与脉冲主路径的输出进行隐藏状态对齐损失计算。这个补偿路径不参与推理,因此不影响最终模型的稀疏性和能效。实践中,研究者仅在全模型(如1.3B模型的48层中)的少数几层(首层、中间层和末层)加入该路径,以平衡训练效率和性能。
第二阶段:两阶段训练框架构建 研究采用了一种高效的训练策略,避免了从头训练脉冲大模型的巨大成本。
蒸馏阶段:采用单阶段知识蒸馏策略,将一个预训练的Mamba2模型作为教师模型,指导SpikingMamba学生模型的训练。训练目标包括两部分:输出层面的KL散度损失,用于对齐教师和学生模型的预测概率分布;以及通过平滑梯度补偿路径实现的隐藏状态对齐损失,用于保证中间层表征的保真度。训练数据使用了多个指令微调数据集(如GenQA、InfinityInstruct、OpenHermes 2.5),仅使用7B tokens进行一个epoch的监督微调。
对齐阶段(可选强化学习):为了进一步提升模型的零样本推理和对齐能力,研究探索了对蒸馏后的SpikingMamba应用强化学习方法,如直接偏好优化(Direct Preference Optimization, DPO)及其变体KTO。这一阶段利用偏好数据(如UltraFeedback)进行微调,耗时极短(少于10个GPU小时),旨在提升模型在特定任务上的表现和与人类偏好的对齐度。
第三阶段:综合实验验证与分析 研究者设计了一系列实验,使用开源的语言模型评估工具,在多个标准常识推理任务(如BoolQ、PIQA、HellaSwag、Winogrande、ARC-Easy/Challenge)上评估模型的零样本准确率,并在语言建模数据集(如WikiText-2、C4、PTB)上评估困惑度。能效分析则通过理论建模,对比了原始Mamba2与不同配置SpikingMamba在45nm工艺下的能耗比。此外,还进行了详尽的消融研究,以验证SI-LIF神经元、平滑梯度补偿路径等关键组件的有效性。
三、主要研究结果
实验结果全面支持了SpikingMamba设计的有效性和高效性。
性能表现:在1.3B参数规模下,SpikingMamba在零样本任务上的平均准确率达到54.62%,与原始Mamba2(59.40%)的差距仅为4.78%。这一表现显著优于其他需要更多训练资源的1位量化Mamba方法(如bi-mamba),也超过了参数量大得多的7B参数SpikeLLM模型(后者相比其教师模型性能下降11.57%)。在困惑度指标上,SpikingMamba也优于从头训练的脉冲语言模型如SpikeGPT和SpikingSSMs。
能效提升:能效分析表明,SpikingMamba在推理时能带来显著的节能效果。在1.3B模型上,相比原始Mamba2,SpikingMamba实现了约4.76倍的能效收益(能耗比E_a/E_s)。这种收益主要来源于将输入/输出投影层的密集矩阵乘法转换为了由稀疏脉冲触发的累加操作。能效比随着模型规模增大而提升,因为投影层的计算占比增加。
组件有效性验证:消融研究证实了关键组件的贡献。使用SI-LIF神经元相比仅支持正值的I-LIF神经元,带来了显著的性能提升(平均准确率提升约0.76%),而SI-LIF相比传统的LIF神经元提升则更为巨大(约6.27%)。平滑梯度补偿路径在仅应用于3层的情况下,能为1.3B模型带来平均0.8%的零样本准确率提升。激活分布分析显示,SI-LIF神经元成功地缓解了脉冲模型的输出极化问题,使其分布更接近原始ANN模型。
强化学习兼容性:研究还发现,经过蒸馏的SpikingMamba可以有效地与强化学习对齐方法(DPO/KTO)结合。经过RL微调后,模型性能得到进一步改善(例如使用KTO后,零样本准确率提升至57.17%),将性能差距从4.78%缩小至2.23%。有趣的是,脉冲模型从RL中获得的相对收益似乎高于其教师ANN模型,这表明SNN架构在奖励驱动的适应性调整中可能有其独特的潜力。
四、研究结论与价值
本研究成功提出了SpikingMamba,这是一种将Mamba架构的序列建模效率与脉冲神经网络的事件驱动节能优势相结合的新型大语言模型。通过创新的SI-LIF神经元设计和平滑梯度补偿路径,研究解决了将预训练大模型能力迁移到脉冲架构时面临的信息保留和梯度优化难题。所采用的单阶段蒸馏与强化学习对齐策略,提供了一条高效、低成本构建高性能脉冲大语言模型的可行路径。
该研究的价值主要体现在以下几个方面: 1. 科学价值:首次将脉冲神经网络与先进的Mamba架构系统性地结合用于十亿参数级别的语言建模,证明了通过精心设计的神经元模型和训练策略,可以在脉冲稀疏计算范式下有效保持大语言模型的语义表征和推理能力,为未来能效导向的神经形态计算与人工智能的融合提供了重要理论和技术基础。 2. 应用价值:SpikingMamba展现出的显著能效优势(高达4.76倍),使其在边缘计算、移动设备、物联网等对功耗和延迟有严格限制的场景中具有巨大的应用潜力。它为部署高效、轻量级的大型语言模型开辟了新的可能性。 3. 方法论价值:提出的SI-LIF神经元和平滑梯度补偿路径具有通用性,可被应用于其他希望引入脉冲计算以提升能效的模型架构中。同时,蒸馏结合强化学习的微调范式也为其他模型压缩和效率优化任务提供了参考。
五、研究亮点
本研究的突出亮点包括: 1. 首创性工作:这是首个支持知识蒸馏、具备线性时间复杂度、并在十亿参数尺度上验证有效的脉冲大语言模型,填补了该领域的研究空白。 2. 关键技术创新:SI-LIF神经元设计巧妙,同时解决了极性保留和幅度量化问题;平滑梯度补偿路径构思新颖,以极小的训练开销显著缓解了量化损失。 3. 卓越的性能-能效权衡:在仅牺牲不到5%零样本准确率的情况下,实现了接近5倍的能效提升,这一权衡结果在当前脉冲语言模型研究中处于领先水平。 4. 高效可扩展的训练流程:避免了代价高昂的从头预训练,仅通过相对少量的蒸馏和微调数据及计算资源,即可将预训练模型的能力成功迁移,具备良好的实用性和可扩展性。 5. 全面的实验验证:研究不仅评估了标准任务性能,还进行了深入的能效理论分析、消融研究和激活分布可视化,为结论提供了坚实的数据支撑和深入洞察。
SpikingMamba研究为大语言模型的高效能效部署提供了一个有前景的新方向,其核心思想和技术路径预计将对人工智能硬件协同设计、绿色AI等领域产生积极影响。