学术报告：SpikeGPT——基于脉冲神经网络的生成式预训练语言模型

作者及发表信息

本研究的核心作者包括：
- Rui-Jie Zhu（加州大学圣克鲁兹分校电气与计算机工程系）
- Qihang Zhao（快手）
- Guoqi Li（中国科学院自动化研究所）
- Jason K. Eshraghian（加州大学圣克鲁兹分校电气与计算机工程系）

论文于2024年6月发表在期刊 *Transactions on Machine Learning Research*，并在开放评审平台OpenReview上公开（链接：https://openreview.net/forum?id=gcf1anbl9e）。

研究背景

科学领域与问题背景
随着大语言模型（Large Language Models, LLMs）规模的持续扩大，其计算资源需求呈指数级增长。例如，训练GPT-3消耗了约19万kWh的能量（Brown et al., 2020）。为应对这一挑战，脉冲神经网络（Spiking Neural Networks, SNNs）因其事件驱动的稀疏激活特性，成为降低计算开销的潜在解决方案。SNNs在计算机视觉任务中已表现出竞争力，但在语言生成任务中尚未取得突破，主要受限于以下问题：
1. 训练困难：SNNs的二元脉冲激活导致层间带宽受限，梯度信号稀疏（Eshraghian et al., 2022）。
2. 与自注意力机制的兼容性：传统Transformer的并行输入与SNNs的时序性存在矛盾（Vaswani et al., 2017）。

研究目标
本研究提出SpikeGPT，首次实现基于直接训练的SNN语言生成模型，旨在：
1. 通过改进Transformer模块，将计算复杂度从二次（O(T²)）降至线性（O(T)）。
2. 在保持能量效率的同时，达到与非脉冲模型（ANNs）相当的性能。

研究方法与流程

1. 模型架构设计

SpikeGPT基于Receptance Weighted Key Value (RWKV)语言模型，通过以下创新实现SNN兼容性：
- 时序对齐：将语言序列维度与SNNs的时间维度对齐，省去额外的时序编码器。
- 自回归训练：采用逐时间步的损失计算，避免传统SNNs的脉冲累积输出方式。
- 状态神经元：利用带状态的神经元（如Leaky Integrate-and-Fire, LIF）缓解二元激活的信息损失。

核心模块（见图1）：
1. Spiking RWKV (SRWKV)：替换传统自注意力机制，通过线性复杂度混合token信息。
- 输入：二元脉冲序列，通过MLP层生成Receptance（R）、Key（K）、Value（V）。
- 动态混合：使用可学习的衰减向量（w_d, w_f）控制历史信息权重。
2. Spiking Receptance Feed-Forward Network (SRFFN)：类似GLU的门控结构，增强信息流控制。
- 激活函数：采用ReLU2或脉冲神经元（LIF），支持稀疏整数脉冲。

2. 训练与推理

两阶段训练流程：
1. 预训练：在大规模语料（如OpenWebText2）上采用自回归目标函数（公式9）。
2. 微调：针对下游任务（如文本生成或分类）调整模型头部：
- 自然语言生成（NLG）：延续预训练目标，生成下一个token的概率分布。
- 自然语言理解（NLU）：通过平均池化token嵌入后接分类头（公式10-12）。

硬件兼容性：
- 理论能耗分析显示，SpikeGPT在神经形态硬件上可减少32.2倍操作量（表1）。
- 关键操作（如稀疏加法和整数乘法）适配异步硬件（如Loihi芯片）。

主要结果

1. 语言生成性能

• EnWik8数据集：SpikeGPT-46M（T=3072）测试BPC（Bits Per Character）为1.262，优于LSTM模型（1.670），接近线性Transformer（1.207）（表2）。
  
• WikiText-103：216M参数模型验证集困惑度（PPL）为39.92，虽落后于GPT-2 small（29.41），但能耗显著降低（表3）。
  

2. 文本分类任务

• SST-2情感分析：预训练后的SpikeGPT-216M准确率达88.76%，接近BERT（91.73%），远超SNN基线TextSCNN（75.73%）（表4）。
  
• 计算效率：模型复杂度为O(T·d)，无需额外时序维度，适合长序列处理。
  

3. 消融实验

• SpikeGPT变体对比（表5）：
  
  • 在SRFFN中使用LIF替代ReLU2（SpikeGPT-B）性能相近（测试BPC 1.306 vs. 1.283），表明二元激活对性能影响有限。
    
  • 纯Heaviside函数（无LIF）性能显著下降（BPC 1.403），验证状态神经元的重要性。
    

结论与价值

科学意义：
1. 首个直接训练的SNN语言生成模型，填补了脉冲计算在生成式任务中的空白。
2. 方法论创新：通过RWKV的线性注意力机制与SNNs的时序性结合，解决了自注意力与SNNs的兼容性问题。

应用价值：
- 低能耗部署：在神经形态硬件上实现高效推理，适合边缘设备。
- 扩展性验证：模型参数从46M至1.2B的缩放实验显示符合神经缩放定律（图3），支持更大规模应用。

研究亮点

1. 能量效率：理论分析表明，SpikeGPT在异步硬件上能耗仅为传统GPT的1/33。
   
2. 架构通用性：SRWKV和SRFFN模块可推广至其他时序建模任务。
   
3. 开源贡献：代码公开于GitHub（https://github.com/ridgerchu/spikegpt），推动SNN社区发展。
   

其他发现

• 对抗性鲁棒性：SNNs在文本分类中表现出对对抗攻击的天然鲁棒性（Lv et al., 2023b）。
  
• 未来方向：结合知识蒸馏（如SpikingBERT）可能进一步提升大规模SNN性能（Bal & Sengupta, 2023）。
  

（注：术语对照：BPC=每字符比特数；PPL=困惑度；LIF=泄漏积分发放神经元；GLU=门控线性单元；RWKV=接收加权键值模型）