paper

SpikingBERT: Distilling BERT to Train Spiking Language Models Using Implicit Differentiation

Abstract

大型语言模型（LLM）虽然变得非常强大，但其神经元和突触的数量比人脑少几个数量级。然而，它需要更多的电力/能量才能运行。在这项工作中，我们提出了一种新颖的仿生脉冲语言模型（LM），旨在通过从大脑中的突触信息流中汲取动力来降低传统 LM 的计算成本。在本文中，我们演示了一个框架，该框架利用平衡状态下神经元的平均脉冲率，使用隐式微分技术训练神经形态脉冲 LM，从而克服基于脉冲神经网络 (SNN) 的算法的不可微性问题，而无需使用任何类型的算法。替代梯度。脉冲神经元的稳态收敛还允许我们设计脉冲注意力机制，这对于开发可扩展的脉冲 LM 至关重要。此外，平衡状态下神经元平均脉冲率的收敛被用来开发一种新颖的基于 ANN-SNN 知识蒸馏的技术，其中我们使用预训练的 BERT 模型作为“老师”来训练我们的“学生”脉冲架构。虽然本文提出的主要架构是由 BERT 推动的，但该技术可以扩展到不同类型的LLMs。我们的工作是第一个在 GLUE 基准测试中展示操作脉冲 LM 架构在多个不同任务上的性能的工作

1.Introduction

大型语言模型(llm)由于在各种自然语言处理(NLP)任务中的广泛应用而变得越来越受欢迎。像GPT-3 (Brown et al. 2020)这样的llm显示出了额外的特征，如突生能力(Wei et al. 2022)，这只能在模型大小/计算量增加到一定阈值以上时才能实现。最近，llm的商业部署实现了全球覆盖，并积极影响了现实世界的用户。然而，llm的巨大能力是以巨大的能量消耗为代价的，无论是在计算昂贵的训练阶段还是在推理阶段。llm的特点是具有大量的可训练参数，并且通常非常深。为了减轻llm的操作复杂性，本文旨在从人类大脑中汲取动力。将llm中体现的机制和知识整合到受大脑启发的神经架构设计中，对于创建生物可信的和节能的解决方案具有巨大的前景。

脉冲神经网络 (SNN)（Ghosh-Dastidar 和 Adeli 2009）受到生物学启发，SNN 架构中两个神经元之间的通信以脉冲的形式发生。这种基于稀疏脉冲的信息流可以在神经形态硬件中实现事件驱动的计算和通信，从而显着节省能源（Sengupta 等人，2019）。基于 SNN 的架构还在英特尔 Loihi 2 处理器（Davies 等人，2021）等神经形态硬件上进行了广泛测试，并展示了数量级的能源效率

将脉冲架构扩展到复杂领域（例如 NLP）会带来重大挑战，主要是由于缺乏可扩展且高效的学习算法。任务的复杂性，加上所需模型架构的深度不断增加，使得 BPTT 的实际应用变得不可行。在这项工作中，在坚实的理论基础和经验证据的支持下，我们探索了一个用于训练尖峰 LM 的可扩展框架。我们将脉冲 LM 视为一个动态系统，在给定输入的情况下，它会逐渐收敛到稳态（在 Tconv 时间步长内）。与大多数监督学习算法类似，训练分两个阶段完成，即。 “前进”和“后退”。然而，我们不是通过在操作的时间步上展开计算图来学习（如 BPTT），而是在学习的“前进”阶段利用神经元的平均尖峰率（ASR）收敛到平衡状态。收敛后，我们可以从基础模型导出定点方程，然后对所获得的稳态采用隐式微分来有效地训练模型参数，如下所述

隐微分训练主要用于深度均衡模型(DEQ) (Bai, Kolter和Koltun 2019)。最近，该方法也被用于视觉相关任务中基于卷积的脉冲架构的训练(Xiao et al. 2021)。与BPTT不同，这种方法在训练期间提供了非凡的内存效率，BPTT需要大量内存来存储大型计算图。通过隐式计算梯度，消除了代理梯度方法的必要性，从而规避了脉冲模型的不可微问题。在一定的约束条件下(Bai、Kolter和Koltun 2019)，这种形式的学习类似于生物可行的和基于能量的训练方法，如平衡传播(Scellier和Bengio 2017;Bal和Sengupta 2022)，从而支持学习的神经形态观点。

在本文讨论的基于 Transformer的 LM 架构中，注意力机制是一个重要的构建块。然而，普通的注意力机制本质上本质上是非脉冲的，因为它依赖于查询、键和值组件的实值向量序列。在本文中，我们提出了一种脉冲注意力机制，该机制利用基于脉冲的输入，并在实现模型收敛所需的时间步数（Tconv）上进行操作。平衡状态下神经元 ASR 的收敛使我们能够在尖峰注意力层的 ASR 和普通注意力之间得出紧密的等价关系。

从头开始培训 LM 是一个耗费大量时间和资源的过程。从头开始训练峰值 LM 的额外开销促使我们寻找更有效的方法来训练我们的模型。知识蒸馏（KD）（Hinton、Vinyals 和 Dean 2015）允许更快、更有效地将知识从经过训练的“教师”模型转移到规模可能更小的“学生”模型。在本文中，我们利用脉冲 LM 的稳态 ASR，提出了一种新颖的 ANN-SNN KD 框架，其中涉及“学生”模型的特定中间层的平衡 ASR 和更大预训练目标层的激活“老师”模式。此外，前面讨论的训练方法使得蒸馏过程中模型训练的可行性成为可能。

我们工作的主要贡献如下：

具有 Spiking Attention 的 SpikingBERT：我们提出了一种完全可操作的脉冲 LM，遵循 BERT 的架构（Devlin 等人，2018 年），并根据 GLUE 基准的不同任务（分类和回归）对其进行评估（Wang 等人，2018 年）。我们还提出了一种有效的脉冲注意力机制，其平衡时的 ASR 近似于普通的非脉冲注意力。
Spiking LM 训练：我们从理论上和实证上验证了我们提出的脉冲 LM（包括线性和非线性操作）到平衡状态的收敛性并使用一种基于隐式微分的方法，克服 SNN 训练的不可微分问题，并减少训练期间的内存使用。这种方法能够训练超过现有脉冲模型规模的脉冲 LM，从而允许为复杂任务开发更深层次的模型。
使用平衡状态的 ANN-SNN KD：我们利用收敛后神经元的平衡状态来更多地训练我们的模型有效地使用新颖的 ANN-SNN KD 框架。这允许使用更大的 BERT 模型作为“老师”来开发高效且更小的脉冲“学生”模型。

2.Related Works

2.1 Spiking Architectures

脉冲架构主要研究了基于视觉的数据集，如CIFAR-100、ImageNet等，以及神经形态数据集，如NMNIST、DVS Gesture。然而，对于基于序列的自然语言处理任务，人们做的工作有限。虽然大多数网络都是基于非注意力的浅层模型，但最近(Zhu、Zhao和Eshraghian 2023)探索了使用线性注意力开发类似gpt的模型。GPT (Radford et al. 2018)是一种纯解码器架构，而BERT是一种纯编码器架构，其捕获双向上下文信息的能力使其更适合文本分类问题。与SpikingBERT不同，SpikeGPT模型使用代理梯度来克服不可微问题，并使用BPTT方法进行训练。此外，据我们所知，SpikingBERT是第一个spike LM，在GLUE基准中根据多个不同的任务进行了评估。此外，由于在该方法中纳入了有效的KD，可以在不需要大量参数的情况下提高模型性能。之前有一些关于SNNs中KD的工作，但他们都主要探索了基于BPTT的方法，并且只关注简单的基于视觉的数据集.

2.2 Efficient LMs

鉴于语言模型的规模不断扩大，致力于降低其计算成本和减小尺寸的研究已经引起了广泛关注。 TinyBERT提出从“教师”模型（中间层和预测层）中提取知识。 NAS-BERT使用神经架构搜索进行模型压缩。我们还开展了将 BERT 知识提炼到单层 BiLSTM 网络的工作。此外，人们还致力于量化、剪枝等方法来降低模型复杂度。由于目前还没有合适的神经形态基线，我们将我们提出的模型与现有的标准 NLP 模型和高效的 LM 进行比较

3.Method

在本节中，我们将首先检查我们方法的基本原则。随后，我们将深入研究架构细节，以提供全面的了解。我们深入研究了使用隐式微分训练的尖峰神经网络中 ASR 收敛的理论和经验基础。此外，我们提出了一种创新方法，利用这种融合来设计脉冲注意力机制和新颖的 KD 机制，利用预先训练的 BERT 模型作为“老师”，从而增强学习过程。我们还详细阐述了使用基于隐式微分的技术来训练我们的脉冲架构的框架

3.1 Spiking Neural Networks

所提出的脉冲架构的基本构建块由漏积分和激发（LIF）神经元组成。简单 LIF 神经元的内部动力学类似于生物神经元，如下所示：

\tau_m · \frac{du}{dt} = -(u(t) - u_{rest}) + R · I(t)

u:膜电位
$u_{rest}$ : 静息膜电位
I(t): 时间 t 时的输入电压，由常数 R（代表电阻）缩放；
$\tau_m$ : 时间常数
如果 u > $V_{th}$ ，则通过减去 Vth 来更新神经元电位，并且神经元发出脉冲，即，它发送“1”否则“0”。

提出的 SpikingBERT 模型的高级概述：在学习的“前向”阶段，网络在 Tconv 时间步长上进行模拟，即直到每层神经元的 ASR 收敛到平衡。两个脉冲编码器内部和之间的信息流使用脉冲而不是真实值发生，从而模仿仿生系统中事件驱动的信息流。

第 i 个神经元动力学的合适离散时间表示可以描述如下:

u_i[t+\delta] = \gamma u_i[t] + \sum_j(w_{ij}s_j[t]) + b_i,

s_i[t+1] = S(u_i[t+\delta]),

u_i[t+1] = u_i[t+\delta] - V_{th}s_i[t+1]

$\gamma$ : 与方程 1 中常数 $\tau_m$ 相关的泄漏项。（对于 LIF 神经元，我们保持 $\gamma$ < 1，对于 IF $\gamma$ = 1）；
$s_j$ : 来自第 j 个输入神经元的脉冲；
$w_{ij}$ : 突触前和突触后神经元之间连接的突触权重；
$t+δ$ : 中间时间步长表示，用于确定神经元是否已激发；
$b_i$ : 代表偏置项；
$S$ : 脉冲生成的不可微函数，减法用作重置操作

3.2 Implicit Modelling

隐式建模采用了一种不同的方法，它不显式地从输入中定义模型输出的精确计算。相反，它依赖于对模型施加特定的约束，确保满足这些约束以获得预期的结果。例如，考虑一个用函数h表示的简单模型。为了表述一个输入x∈X，输出z∈Z的显式模型，需要进行以下计算:z = h(x)。然而，为了隐式地表达它，定义了一个函数g: X × Z→ $R^n$ ，使g(x, z) = h(x)−z，目标是找到方程的根:g(x, z) = 0。虽然这个简单的例子演示了代数方程，但这些方法可以扩展到不动点方程，从而为DEQ的发展铺平了道路(Bai, Kolter和Koltun 2019)。

让我们考虑 $z =f_θ(z)$ 形式的定点方程，其中 θ 是参数集。这个定点方程随着时间的推移而收敛，即经过 $T_{conv}$ 时间步长 $z_{T_{conv}} = z(T_{conv} + 1)$ ，从而达到平衡状态。类似地，和之前一样，我们可以形成另一个方程，即 $g_θ(z) = f_θ(z) − z$ 。在这里，我们将定义的损失函数 L 将利用平衡时的 z 值，即 $z_{T_{conv}} = z^*$ 。使用隐式微分 (Bai, Kolter, andKoltun 2019)，可以得出以下关系:

\frac{\partial L(z^*)}{\partial \theta} = \frac{\partial L(z^*)}{\partial z*}(J^{-1}_{g_{\theta}}|z^*) \frac{f_{\theta}(z^*)}{\partial \theta}

$J^{-1}_{g_{\theta}}|z^*$ : $z = z^*$ 时 gθ 的逆雅可比行列式，即处于平衡状态。然而，雅可比行列式的逆计算在计算上非常昂贵。因此，我们利用 Broyden 方法等迭代方法（Bai、Koltun 和 Kolter 2020）并求解由 $(J^{-1}_{g_{\theta}}|z^*)x + (\frac{\partial L(z^*)}{\partial z^*})^T$ = 0 给出的线性方程组。该方程的解为我们提供了 $−\frac{\partial L(z^*)}{\partial z^*}(J^{-1}_{g_{\theta}}|z^*)$ ，我们可以直接通过均衡进行反向传播（Bai、Kolter 和 Koltun 2019）。所提出的脉冲 LM 架构遵循一组类似的方程，将在下一节中进行描述。由于梯度是使用收敛稳态上的隐式微分来计算的，因此我们避免了脉冲函数的不可微分问题。此外，通过仅在平衡状态计算梯度，不需要存储中间隐藏状态。与 BPTT 相比，这一特性增强了该方法的可扩展性和内存效率。

3.3 Architecture

所提出的脉冲 LM 的高级构建块由脉冲编码器 (SE) 层组成，可以将其视为类似于 Transformer 架构中的各个编码器层。 SE 层中的层内和层间通信都在“前进”阶段的每个时间步都使用脉冲进行，脉冲LIF 神经元是其设计中的基本单元。如图 1 所示，每个 SE 层都包含一个脉冲注意力层，后面是全连接层（有些包括类似于 BERT 中的跳跃连接），即。中间层 1 (IL-1)、中间层 2 (IL-2) 和输出层，所有这些都使用脉冲运行。所提出的网络结构由 N 个堆叠的 SE 层组成，类似于 BERT 的结构。输入嵌入由 LIF 神经元层处理，生成通过模型传播的脉冲。在SE的一些内部层中，我们使用层归一化（遵循BERT）。附录 B C 中添加了规范化的效果和其他架构细节。

从生物学的角度来看，人类大脑中存在反馈连接，而且在某些情况下具有循环连接的浅层网络(Kubilius等人2019)显示出与深层网络结构相当或更好的性能。连接(F)从最后一个SE层的输出添加到第一个，以引入反馈。反馈连接是一个可选组件，增加了模型的生物可信性。随后发展的稳态ASR方程的一般形式，可以无缝地应用于包括反馈连接和无反馈的模型。与基于视觉的任务不同，在使用GLUE基准的实验中，与无反馈场景相比，反馈没有显著提高性能。然而，我们仍然在理论层面上进行探索，以保持与之前工作的一致性(Xiao等人，2021)，并鼓励未来在其他领域对反馈支持的snn进行研究

与基于普通transformer的架构类似，输入序列直接输入到模型中。该模型在“前进”阶段在 $T_{conv}$ 时间步长内收敛，以达到平衡状态。如前所述，模型的脉冲神经元具有各自的膜电位u，在收敛过程中的每个时间步都会更新。在时间t+1时， $u_1$ 的膜电位（即第一个 SE 层的输入）可以给出为:

u_1[t+1] = \gamma u_1[t] + Fs(N,out)[t]+W_0(x)+b_1-V_{th}s_1[t+1]

$W_0$ 提供长度为 $N_s$ 的输入序列 x 的嵌入，并生成向量序列 y ∈ $R^{N_s × Demb}$ （Demb 是编码维度
F 是反馈连接的权重（如果包括反馈）
$b_1$ : 偏差
$s_{(N,out)}[t]$ : 前一时间步中第 N 个 SE 层生成的脉冲

i > 1 层的膜电位可以简单地表示为：

u_i[t+1] = \gamma u_i[t] + W_{(i-1)}(s_{(i-1)}[t+1]) + b_i -V_{th}s_i[t+1]

$W_{(i−1)}$ : 是在前一层的一组脉冲上定义的操作（由每个单独层制定）

所描述的 LIF 神经元使用按照方程式2生成的脉冲传播信息。

i 层神经元的平均脉冲率 (ASR) 可以定义为： $a_i[t] =\frac{\sum^t_{\tau=1}\gamma^{t-\tau}s_i[\tau]}{\sum^t_{\tau =1}\gamma^{t-\tau}}$ 。给定 $W_{(i−1)}$ 是线性运算，使用方程5并随着时间的推移执行加权（如果 γ < 1）平均值（u[0] = 0，s[0] = 0），我们得到:

a_i[t+1] = \frac{1}{V_{th}}(W_{(i-1)}a_{(i-1)}[t+1] + b_i - \frac{u_i[t+1]}{\sum_{i=0}^{t}\gamma^i})

由于 $a_i[t]$ 代表ASR，因此其值被限制在[0,1]内。继之前关于平衡时隐式微分的工作（Xiao et al. 2021）之后，当输入的平均值收敛到平衡 $\bar{x}[t] → x^＊$ 时，脉冲架构中各层的 ASR（方程 6）收敛到平衡点： $a_i[t]→a^*_i$ （在 LIF 神经元的情况下具有有界随机误差）。在平衡状态下，第 i 层的 ASR $a^*_i$ 满足：

a_i^* = \sigma(\frac{1}{V_{th}}(W_{(i-1)}(a^*_{i-1}+b_i)))

裁剪函数 σ(x) 将值限制在 [0,1] 范围内

图 2：将 SST-2 数据集的随机采样输入传递到 $SpikingBERT_4$ 后获得的结果。
(a) 图表显示 SE 层中不同子层的 ASR 平均值（相对于神经元数量）与操作时间步长的关系。
(b) 左侧的 y 轴描绘了随机选择的脉冲注意力层的 ASR 平均值（除神经元数量之外）。沿着右侧 y 轴，显示了所选脉冲注意力层的稳态方程的输出与计算出的 ASR 之间的“差异范数”。沿 x 轴显示用于收敛的时间步长。

与线性操作一样，具有非线性操作（例如脉冲注意力）的层也会收敛到稳态 ASR，正如本文通过经验验证的那样（图 2a）。

像Eqn. 7这样的稳态方程在训练过程中被利用。因此，对于脉冲注意力层，我们制定了平衡时的替代稳态方程，如下所示：

a^*_{(attn)} = \sigma(\frac{1}{V_{th}}(Attn(a^*_x,a^*_k,a^*_v)+b_{(attn)}))

$a^*_x$ 是用于形成Query的层的ASR
$a^*_k$ 是Key的ASR
$a_v^*$ 是Value的ASR

脉冲注意力机制的操作细节以及定义方程的经验收敛和合理性将在下一小节中讨论

因此，在模型收敛到平衡后，底层 SNN 的稳态 ASR 的动态可以映射到替代非脉冲架构，其中每层的输入和输出都是相应的 ASRs( $a^*_i$ )代理网络中各个层的操作由前面描述的稳态方程给出，并且可以简化为形式， $a^*_i = l_i(a^*_j,...)$ ，其中， $l_is$ 是对应的稳态方程到每一层，如Eqn.7和8。与每个层（ $l_i$ ）相关的参数（ $a^*_j$ ，…）根据具体操作来定义。如果我们使用反馈连接，第一层的定点方程的形式为 $a_1^* = l_1(l_M ◦ · · · ◦ l_2(a_1^*), x^*)$ 其中 $l_1(a, x) = \sigma(\frac{1}{V_{th}}(Fa + W_0(x) + b_1))$ 其中 M 是各个层的总数。

对于文本分类的任务，我们使用网络的最后一层，即第 N 个编码器层的输出，以 $a^*_{N,out}$ 形式给出，作为线性分类函数的输入。模拟 $T_{conv}$ 时间步长的网络，我们可以计算 $a_{N,out}[T] =\frac{\sum_t(s_{N,out}[t])}{T}$ （为了简化演示，γ = 1），我们可以将其用作 $a^*_{N,out}$ 。此外，由于代理网络仅使用 ASR 捕获平衡时的行为，因此我们可以简单地利用等式3执行反向传播来训练权重。因此，我们不使用 BPTT 来训练底层脉冲架构，而是使用简单的反向传播，仅使用神经元的平衡状态 ASR 来训练脉冲 LM 的权重。