[论文解读] SpikingBERT: Distilling BERT to Train Spiking Language Models Using Implicit Differentiation

Malyaban Bal, Abhronil Sengupta|arXiv (Cornell University)|Aug 21, 2023

Advanced Memory and Neural ComputingEngineering被引用 3

一句话总结

SpikingBERT 提出了一种受生物启发的、能效更高的脉冲语言模型，通过隐式微分训练以绕过脉冲神经网络的不可微性，实现无需替代梯度的端到端训练。它利用一种新颖的 ANN-SNN 知识蒸馏框架和基于稳态脉冲发放率的脉冲注意力机制，在 GLUE 基准测试中取得了具有竞争力的性能表现。

ABSTRACT

Large language Models (LLMs), though growing exceedingly powerful, comprises of orders of magnitude less neurons and synapses than the human brain. However, it requires significantly more power/energy to operate. In this work, we propose a novel bio-inspired spiking language model (LM) which aims to reduce the computational cost of conventional LMs by drawing motivation from the synaptic information flow in the brain. In this paper, we demonstrate a framework that leverages the average spiking rate of neurons at equilibrium to train a neuromorphic spiking LM using implicit differentiation technique, thereby overcoming the non-differentiability problem of spiking neural network (SNN) based algorithms without using any type of surrogate gradient. The steady-state convergence of the spiking neurons also allows us to design a spiking attention mechanism, which is critical in developing a scalable spiking LM. Moreover, the convergence of average spiking rate of neurons at equilibrium is utilized to develop a novel ANN-SNN knowledge distillation based technique wherein we use a pre-trained BERT model as "teacher" to train our "student" spiking architecture. While the primary architecture proposed in this paper is motivated by BERT, the technique can be potentially extended to different kinds of LLMs. Our work is the first one to demonstrate the performance of an operational spiking LM architecture on multiple different tasks in the GLUE benchmark.

研究动机与目标

• 开发一种生物上合理、低功耗的脉冲语言模型（SNN），在保持性能的同时降低大语言模型（LLMs）的能耗。
• 通过使用隐式微分而非替代梯度，克服 SNN 的不可微性，实现端到端训练。
• 设计一种基于稳态平均脉冲发放率（ASR）的脉冲注意力机制，以兼容基于 Transformer 的架构。
• 利用平衡态 ASR 和中间层激活，实现从预训练 BERT '教师' 模型到脉冲 '学生' 模型的高效知识蒸馏。
• 首次在 GLUE 基准上实现可运行的脉冲语言模型，展示了在能效方面的可行性与性能提升。

提出的方法

• 在脉冲神经元在 $T_{\text{conv}}$ 时间步内收敛至稳态的固定点方程上应用隐式微分，避免使用时间反向传播（BPTT）和替代梯度。
• 提出一种脉冲注意力机制，其中查询、键和值操作基于脉冲输入计算，并利用平衡态下的平均脉冲发放率（ASR），确保与标准注意力机制等价。
• 采用一种新颖的 ANN-SNN 知识蒸馏框架，将脉冲学生模型中的中间层 ASR 与预训练 BERT '教师' 模型的激活模式对齐。
• 实施多阶段蒸馏：通用蒸馏、任务特定蒸馏（如在 MNLI 上，随后在 QNLI 上），以及最终输出层蒸馏，以提升准确率。
• 使用归一化操作（Norm\#OPS）和推导出的能效因子 $e = \big(\frac{1}{5.1} \times \text{Norm\\#OPS}\big)^{-1}$ 来量化能效节省，其中 ACC 操作仅消耗 0.9pJ，而 MAC 操作消耗 4.6pJ。
• 对 $V_{\text{th}}$ 和 $T_{\text{conv}}$ 进行消融研究，分析在 SST-2 上功耗、收敛时间与准确率之间的权衡。

Figure 1: High-level overview of the SpikingBERT model. During the “forward” phase of learning, the network is simulated over $T_{conv}$ time steps, i.e., until the ASR of the neurons of each layer converges to an equilibrium. Information flow both within and between two spiking encoders occur using

实验结果

研究问题

• RQ1 能否在不使用替代梯度或 BPTT 的情况下，通过隐式微分处理稳态脉冲动力学，有效训练基于脉冲 Transformer 的语言模型？
• RQ2 如何设计一种脉冲注意力机制，使其在基于脉冲的事件驱动框架中保持标准自注意力的表征能力？
• RQ3 能否通过基于 ASR 的蒸馏方法，将预训练 BERT 模型的知识高效迁移至脉冲学生模型？这种迁移是否能提升性能？
• RQ4 SpikingBERT 的能效-准确率权衡如何？它能否在保持具有竞争力的准确率的同时，实现相比非脉冲模型显著的能效提升？
• RQ5 超参数如 $T_{\text{conv}}$ 和 $V_{\text{th}}$ 如何影响脉冲语言模型中的能耗、收敛时间与模型准确率？

主要发现

• SpikingBERT 在 GLUE 基准测试中实现了具有竞争力的性能表现，当省略知识蒸馏时，准确率下降 4–5%，证明了所提蒸馏框架的必要性。
• 在 $T_{\text{conv}} = 16$ 时，SpikingBERT 在 SST-2 上达到最大准确率的 98%，同时相比同尺寸的非脉冲模型实现近两倍的能效提升（$e$）。
• 提高膜电位阈值 ($V_{\text{th}}$) 可降低平均脉冲发放率（ASR）和瞬时功耗，实现更低功耗运行，适用于边缘设备。
• 能效因子 $e$ 的推导公式为 $e = \big(\frac{1}{5.1} \times \text{Norm\\#OPS}\big)^{-1}$，表明 SNN 中的 ACC 操作比 MAC 操作能效高出 5 倍以上。
• 任务特定知识蒸馏（如从 MNLI 到 QNLI）相比通用蒸馏表现更优，在下游任务上蒸馏可带来准确率提升。
• 脉冲注意力机制通过将脉冲神经元的稳态 ASR 与注意力输出对齐，成功模拟了标准注意力机制，实现在平衡态下的功能等价。

Figure 2: Results obtained after passing a randomly sampled input from SST-2 dataset through SpikingBERT 4 . (a) Graph showing mean (over number of neurons) of the ASR of different sub-layers in an SE layer against the operating time steps. (b) The y-axis on the left depicts mean (over number of neu