2025-11-17T01:31:12.954580

Large Language Models Inference Engines based on Spiking Neural Networks

Balaji, Madireddy, Balaprakash
Foundational models based on the transformer architecture are currently the state-of-the-art in general language modeling, as well as in scientific areas such as material science and climate. However, training and deploying these models is computationally challenging as the time and space complexity has a quadratic relation to the input sequence length. Several efforts exploring efficient computational paradigms and model architectures to address these limitations have been made. In this work, we explore spiking neural networks (SNNs) to design transformer models. A challenge in training large-scale SNNs, using existing surrogate learning methods is inefficient and time-consuming. On the other hand, techniques to convert existing transformer-based models to their SNN equivalent are not scalable, as achieving optimal performance comes at the cost of a large number of spike time-steps, i.e. increased latency. To address this, we propose NeurTransformer, a methodology for designing transformer-based SNN for inference using a supervised fine-tuning approach with existing conversion methods. The proposed methodology works by: (1) replacing the self-attention mechanism with a spike-based self-attention (SSA), (2) converting the feed-forward block of the trained transformer model to its equivalent SNN, and (3) fine-tuning the SSA block using SNN-based surrogate learning algorithms. We benchmark the proposed methodology and demonstrate its accuracy and scalability using three variants of the GPT-2 model of increasing model size. We observe that the converted GPT-2 small models demonstrate a 5-12% loss in cosine similarity and a 9.7% reduction in perplexity. Finally, we demonstrate the energy efficiency of the SSA block compared to the ASA block and show between 64.71% and 85.28% reductions in estimated energy consumption when implementing the self-attention mechanism on a digital hardware.
academic

Large Language Models Inference Engines based on Spiking Neural Networks

基本信息

  • 论文ID: 2510.00133
  • 标题: Large Language Models Inference Engines based on Spiking Neural Networks
  • 作者: Adarsha Balaji (Argonne National Laboratory), Sandeep Madireddy (Argonne National Laboratory), Prasanna Balaprakash (Oak Ridge National Laboratory)
  • 分类: cs.LG (Machine Learning)
  • 发表时间: 2025年10月14日 (arXiv预印本)
  • 论文链接: https://arxiv.org/abs/2510.00133v3

摘要

基于Transformer架构的基础模型在通用语言建模以及材料科学、气候科学等科学领域都是当前的最先进技术。然而，这些模型的训练和部署在计算上极具挑战性，因为其时间和空间复杂度与输入序列长度呈二次关系。本文探索使用脉冲神经网络(SNNs)来设计Transformer模型。现有代理学习方法训练大规模SNNs效率低下且耗时，而将现有Transformer模型转换为等效SNN的技术缺乏可扩展性。为此，作者提出了NeuTransformer方法，通过监督微调结合现有转换方法来设计基于SNN的Transformer推理引擎。该方法包括：(1)用基于脉冲的自注意力(SSA)替换自注意力机制，(2)将前馈块转换为等效SNN，(3)使用SNN代理学习算法微调SSA块。实验表明，转换后的GPT-2小模型在余弦相似度上损失5-12%，困惑度降低9.7%，SSA块相比ASA块实现64.71%-85.28%的能耗降低。

研究背景与动机

核心问题

  1. 计算复杂度问题：Transformer模型的自注意力机制具有O(n²)的时间和空间复杂度，其中n为序列长度，这使得处理长序列时计算和内存需求急剧增加。
  2. 能耗问题：传统Transformer模型需要昂贵的GPU或定制加速器进行训练和推理，能耗巨大。
  3. SNN训练困难：现有的SNN训练方法存在两个主要局限：
    • 直接训练大规模SNN使用反向传播学习规则效率低下
    • ANN-SNN转换方法需要大量脉冲时间步长才能达到最优性能，导致推理延迟增加

研究动机

作者希望利用脉冲神经网络的生物启发特性和事件驱动计算优势，设计能够在神经形态硬件上高效执行的Transformer模型，从而实现数据高效、能耗低、资源节约的大型语言模型推理。

核心贡献

  1. 提出NeuTransformer方法：一种从训练好的Transformer模型设计基于SNN的Transformer的方法，结合监督微调来提升模型性能。
  2. 设计稀疏脉冲计算的自注意力机制：用稀疏的基于脉冲的计算替代传统自注意力中能耗和延迟较高的矩阵乘法和softmax操作。
  3. 实现大规模SNN-based LLM：成功将GPT-2及其变体转换为SNN版本，据作者所知，GPT-2 Large是目前参数量最大的基于SNN的Transformer模型。
  4. 全面性能评估：从应用准确性、余弦相似度、困惑度、比特每字节等多个维度评估模型性能，并分析能耗和吞吐量表现。

方法详解

任务定义

将预训练的Transformer模型转换为等效的脉冲神经网络版本，同时保持可接受的性能损失，并实现显著的能耗降低和硬件效率提升。

模型架构

1. 脉冲神经元模型

使用具有固定阈值和可调膜电位衰减特性的积分发放(IF)神经元：

S(t) = {
  1,   if Vmem ≥ 1
  -1,  if Vmem ≤ -1  
  0,   otherwise
}

其中Vmem是膜电位，S(t)是神经元输出的脉冲激活。

2. 脉冲自注意力(SSA)机制

传统自注意力机制：

ASA(Q,K,V) = softmax(Q·K^T)V

脉冲自注意力机制：

AttentionScore(AS) = LIF((Q⊗K^T)_Columnwise)
SSA(Q,K,V) = (AS ⊗ V)

关键创新点：

  • 用AND操作和累加器替代N位矩阵乘法
  • 用列级Hadamard乘积替代点积操作
  • 用LIF神经元激活替代softmax函数

3. 脉冲前馈层转换

基于ANN-SNN转换原理，将ReLU激活的前馈层转换为IF神经元：

  • ReLU函数：ReLU(y) = max(0, y)
  • IF神经元：τm ∂Vmem/∂t = -Vmem(t) + R*I(t)

权重归一化：

s^l_norm = max(a^l)
W̃^l ← W^l / s^l_norm

NeuTransformer三步转换流程

  1. 替换自注意力块：将ASA替换为SSA，保留训练好的权重
  2. 转换前馈块：将ReLU/GeLU前馈层转换为SNN等效版本
  3. 微调SSA块：使用代理梯度学习算法微调SSA块权重

微调目标函数

最小化ASA和SSA注意力分数之间的均方误差：

Σ(i=1 to d_model) (ASA_as - SSA_as)²

实验设置

数据集

  • Shakespeare数据集：包含40,000行莎士比亚戏剧文本
  • OpenWebText数据集：OpenAI WebText数据集的开源复现版本

模型规模

  • GPT-2 Small：117M参数
  • GPT-2 Medium：345M参数
  • GPT-2 Large：763M参数

评价指标

  • 字符准确率：ANN与SNN生成字符的逐字符比较
  • 余弦相似度：多维空间中两个非零向量夹角的余弦值
  • 困惑度(Perplexity)：衡量语言模型质量的指标
  • 比特每字节(BpB)：预测下一个token所需的平均比特数

硬件平台

  • 计算资源：Argonne国家实验室LCRC的Swing HPC集群
  • 配置：6个节点，每节点2×AMD EPYC 7742处理器，8×NVIDIA A100 GPU
  • 评估平台：NVIDIA A100 GPU和Graphcore IPU平台

实验结果

主要性能结果

模型参数量余弦相似度字符准确率ANN困惑度SNN困惑度
GPT-2-Small117M0.8884.9%17.1121.81
GPT-2-Medium345M0.8375.4%14.4319.73
GPT-2-Large763M0.7471.8%12.6718.10

能耗分析结果

SSA块相比ASA块的估算能耗降低：

  • GPT-2 Small：85.28%
  • GPT-2 Medium：85.22%
  • GPT-2 Large：64.71%

吞吐量评估

在Graphcore平台上，SNN版本在多数配置下显示出优于基线ANN的吞吐量表现，特别是在处理不规则和稀疏数据访问时优势明显。

关键发现

  1. 规模效应：随着模型规模增大，SNN版本的性能损失逐渐增加
  2. 能效优势：所有规模的模型都实现了显著的能耗降低
  3. 硬件适配性：SNN在MIMD处理架构上表现出色，特别适合稀疏脉冲工作负载

相关工作

SNN-based Transformer研究

  • Spikformer (Li et al., 2024)：首次在视觉任务上实现基于SNN的Transformer
  • Spikingformer (Zhou et al., 2023)：改进版本，使用预激活快捷连接
  • SGLFormer (Zhang et al., 2024)：结合局部和全局Transformer块

ANN-SNN转换方法

  • 早期工作主要针对视觉任务的小规模网络
  • 现有方法在大规模语言模型上的应用受限

本文优势

相比现有方法，NeuTransformer避免了从头训练SNN的计算开销，通过转换+微调策略实现了大规模SNN-based LLM的构建。

结论与讨论

主要结论

  1. 可行性验证：成功将GPT-2系列模型转换为SNN版本，证明了大规模SNN-based LLM的可行性
  2. 性能权衡：在可接受的性能损失下实现了显著的能耗降低
  3. 规模限制：发现当模型参数超过300M时，性能退化超出可接受阈值

局限性

  1. 规模瓶颈：大规模模型的性能退化问题仍需解决
  2. 转换精度：前馈块的不完美转换影响整体性能
  3. 微调限制：代理梯度学习在深度SNN上的局限性

未来方向

  1. 改进大规模SNN的训练和转换方法
  2. 优化代理梯度学习算法
  3. 探索更高效的脉冲编码和解码策略
  4. 在实际神经形态硬件上验证性能

深度评价

优点

  1. 创新性强：首次实现大规模SNN-based语言模型，技术路线新颖
  2. 实用价值高：显著的能耗降低对实际应用具有重要意义
  3. 评估全面：从多个维度评估模型性能，实验设计严谨
  4. 写作清晰：技术描述详细，方法论述清楚

不足

  1. 规模限制明显：大模型性能退化严重，限制了方法的适用性
  2. 理论分析不足：缺乏对性能退化原因的深入理论分析
  3. 硬件验证有限：主要基于估算，缺乏在真实神经形态硬件上的验证
  4. 对比实验不够：与其他SNN-based方法的直接对比较少

影响力

  1. 学术贡献：为SNN在大规模语言模型中的应用开辟了新方向
  2. 实用前景：为低功耗AI推理提供了新的技术路径
  3. 可复现性：方法描述详细，具有较好的可复现性

适用场景

  1. 边缘计算：资源受限的移动设备和IoT场景
  2. 低功耗推理：对能耗敏感的应用场景
  3. 神经形态计算：专门的神经形态硬件平台

参考文献

论文引用了该领域的重要工作，包括：

  • Transformer原始论文 (Vaswani et al., 2017)
  • SNN转换方法的经典工作 (Rueckauer et al., 2016; Diehl & Cook, 2015)
  • 近期SNN-based Transformer研究 (Li et al., 2024; Zhou et al., 2023)
  • 代理梯度学习方法 (Eshraghian et al., 2023)

总结：本文提出了一种将大规模语言模型转换为脉冲神经网络版本的创新方法，在保持可接受性能的同时实现了显著的能耗降低。尽管存在规模限制等问题，但为低功耗AI推理和神经形态计算领域提供了有价值的技术贡献。