paper

SPIKEBERT: A LANGUAGE SPIKFORMER LEARNED FROM BERT WITH KNOWLEDGE


abstract

脉冲神经网络 (SNN) 提供了一种以更节能的方式实现深度神经网络的有前途的途径。然而，现有SNN用于语言任务的网络架构仍然过于简单且相对浅层，深层架构尚未得到充分探索，导致与BERT等主流基于Transformer的网络相比存在显着的性能差距。 为此，我们改进了最近提出的脉冲变压器（即 Spikformer），使其能够处理语言任务，并提出了一种两阶段知识蒸馏方法来训练它，该方法将通过从 BERT 中蒸馏知识来进行预训练与大量未标记文本的集合，并通过知识蒸馏对特定于任务的实例进行微调，再次从 BERT 中对相同的训练示例进行微调。通过广泛的实验，我们表明，使用我们的方法训练的模型（称为 SpikeBERT）超越了最先进的 SNN，甚至在英语和中文文本分类任务上取得了与 BERT 相当的结果，而且能耗要低得多。我们的代码可在 https://github.com/Lvchangze/SpikeBERT 获取。


1.Introduction

  • Spikformer：CV

  • SpikBERT：NLP



如图1所示，考虑到文本数据的离散性，我们改进了Spikformer的架构，使其适合语言任务。我们用语言友好的模块替换了最初为图像处理设计的某些模块。

一般来说，更深的 ANN 模型通常可以带来更好的性能，并且增加 ANN 的深度可以从输入数据中提取更复杂和抽象的特征。然而，方等人。 （2020b）表明，使用代理梯度（参见第 2.1 节）直接通过时间反向传播训练的深度 SNN 会因“自累积动态”而遭受梯度消失或爆炸的问题。因此，我们建议使用知识蒸馏来训练语言Spikformer，这样脉冲神经网络中代理梯度的偏差就不会快速累积。

然而，如何从 ANN 中提取知识到 SNN 是一个巨大的挑战，因为 ANN 中的特征是浮点格式，而 SNN 中的特征是脉冲序列，甚至还涉及额外的维度 T（时间步）。我们发现这个问题可以通过引入外部模块来在训练时对齐特征来解决

受到广泛使用的“训练前微调”配方的启发，我们提出了一种两阶段知识蒸馏策略：

  • 第一阶段：我们选择大型预训练模型的代表例子BERT作为教师模型，SpikeBERT作为学生模型。我们利用大量未标记文本来对齐嵌入层和隐藏层中两个模型产生的特征。
  • 第二阶段：我们作为教师采用针对特定任务数据集进行微调的 BERT，并作为学生采用第一个预训练阶段后的模型。我们首先增加特定于任务的数据集的训练数据，然后使用教师模型预测的逻辑来进一步指导学生模型。经过两阶段的知识蒸馏，可以通过从 BERT 中蒸馏知识来构建脉冲语言模型，称为 SpikeBERT

本文改进了Spikformer的架构，用于语言处理，并提出了一个两阶段的"预训练+特定任务"知识蒸馏训练方法，在此方法中，spikebert在大量未标记文本集上进行预训练，然后通过从bert中提取特征提取和预测能力的知识，在特定任务数据集上进一步微调。


2.Related Work

2.1 SPIKING NEURAL NETWORKS

与传统的人工神经网络不同，脉冲神经网络利用离散的脉冲序列而不是连续的十进制值来计算和传输信息。由于脉冲的不可微性，训练 SNN 在过去二十年一直是一个巨大的挑战。目前解决这个问题的主流方法有两种：

  • ANN-to-SNN Conversion：ANN-SNN转换方法旨在转换良好的重量通过用脉冲神经元层替换激活函数并添加权重归一化（Diehl 等人，2016）和阈值约束（Hu 等人，2018）等缩放规则，将 ANN 训练为其 SNN 对应物。在转换过程中，这种方法遭受了大量的时间步骤。

  • Backpropagation with Surrogate Gradients：另一种流行的方法是在误差反向传播期间引入代理梯度，使整个过程可微。人们提出了多种代理梯度函数，包括Sigmoid代理函数、Fast-Sigmoid、ATan等。随着时间的反向传播（BPTT）是直接训练 SNN 的最流行的方法之一，它将传统的反向传播算法应用于展开的计算图。我们选择带有代理梯度的 BPTT 作为我们的训练方法。


---------- ANN-to-SNN Conversion·Example·BEGIN ----------

在ANN中,每个神经元的输出是一个连续的值,通常是通过一个激活函数如ReLU或Sigmoid函数计算得到的。在SNN中,神经元的输出是脉冲(或称为尖峰),它们是在神经元的膜电位超过某个阈值时发生的。

假设我们有一个已经训练好的ANN,它的权重矩阵是WW,激活函数是ff。

  1. 激活函数替换:对于ANN中的每一个神经元,我们有:

y=f(∑iWixi+b)y = f(\sum_{i} W_{i} x_{i} + b)

在SNN中,我们用脉冲神经元模型替换ff,这里xix_{i}是输入,yy是输出,bb是偏置。

  1. 权重归一化:在转换为SNN后,权重可能会进行归一化:

W′=Wmax⁡(∣W∣)W' = \frac{W}{\max(|W|)}

这里max⁡(∣W∣)\max(|W|)是权重矩阵中绝对值最大的权重,用于将所有权重缩放到[-1, 1]的范围内。

  1. 阈值调整:设VthreshV_{\text{thresh}}是SNN中神经元的阈值,可以根据ANN中的激活函数最大值进行调整:

Vthresh=max⁡(f)V_{\text{thresh}} = \max(f)

这意味着,一旦膜电位累积到VthreshV_{\text{thresh}},神经元就会发放一个脉冲,并重置电位。

  1. 缩放规则:为了确保SNN的动态行为与ANN相匹配,可能需要进一步的缩放规则来调整神经元的行为。这可能包括调整膜电位的衰减或增长率,以匹配ANN中激活函数的曲线。

  2. 时间步骤的处理:在SNN中,我们需要在多个时间步骤上模拟神经元的动态行为。如果我们使用时间步骤tt,则神经元的膜电位V(t)V(t)可以根据以下规则更新:

V(t+1)=λV(t)+W′x(t)−VthreshS(t)V(t+1) = \lambda V(t) + W' x(t) - V_{\text{thresh}} S(t)

其中S(t)S(t)表示在时间步骤tt神经元是否发放脉冲,如果发放了脉冲,则为1,否则为0。λ\lambda是一个介于0和1之间的因子,表示电位的衰减。

在这个过程中,神经元只有在其膜电位超过 VthreshV_{\text{thresh}} 时才发放脉冲。通过这种方式,一个训练有素的ANN可以被转换为SNN,其行为在很大程度上保留了原始网络的特性,但现在它能够以时间序列的形式处理信息。

---------- ANN-to-SNN Conversion·END ----------


---------- Backpropagation with Surrogate Gradients·BEGIN ----------

在SNN中,神经元在时间步tt 的膜电位 V(t)V(t) 可以根据输入 I(t)I(t) 和前一个时间步的电位 V(t−1)V(t-1) 计算得到:

V(t)=V(t−1)+I(t)−S(t−1)⋅VresetV(t) = V(t-1) + I(t) - S(t-1) \cdot V_{\text{reset}}

  • S(t−1)S(t-1) 表示在时间t−1t-1 的发放情况,如果神经元发放了脉冲,则 S(t−1)=1S(t-1) = 1,否则为00
  • VresetV_{\text{reset}} 是一个常数,表示电位重置的值。

神经元是否发放脉冲由一个阶跃函数 HH 决定,它是一个不可微分的函数:

S(t)=H(V(t)−Vthresh)S(t) = H(V(t) - V_{\text{thresh}})

这里 VthreshV_{\text{thresh}} 是电位的阈值。但是,由于 HH 不可微分,我们无法直接用它来计算反向传播的梯度。


使用代理梯度方法,我们定义一个可微分的代理函数 GG,例如,对于阶跃函数,代理函数可能是一个Sigmoid函数:

G(x)=11+e−βxG(x) = \frac{1}{1 + e^{-\beta x}}

在反向传播过程中,我们不直接计算 HH 的梯度,而是计算 GG 的梯度。即使 GG 本身并不是网络前向传播时使用的函数,它仍然提供了关于阈值的有用梯度信息。

考虑到一个简化的损失函数 EE,它直接依赖于时间步tt 的脉冲输出 S(t)S(t),其对 V(t)V(t) 的梯度可以表示为:

∂E∂V(t)=∂E∂S(t)⋅∂S(t)∂V(t)≈∂E∂S(t)⋅G′(V(t)−Vthresh)\frac{\partial E}{\partial V(t)} = \frac{\partial E}{\partial S(t)} \cdot \frac{\partial S(t)}{\partial V(t)} \approx \frac{\partial E}{\partial S(t)} \cdot G'(V(t) - V_{\text{thresh}})

其中,G′(x)G'(x) 是 G(x)G(x) 相对于 xx 的导数。

最后,梯度 ∂E∂V(t)\frac{\partial E}{\partial V(t)} 会被用于通过梯度下降方法更新网络权重:

Wij(new)=Wij(old)−η∂E∂V(t)W_{ij}^{\text{(new)}} = W_{ij}^{\text{(old)}} - \eta \frac{\partial E}{\partial V(t)}

这里 WijW_{ij} 表示连接 ii 和 jj 神经元的权重,而 η\eta 是学习率。

通过以上步骤,带有代理梯度的BPTT使得原本不可微的SNN能够通过标准的梯度下降方法来训练,这是处理不可微操作的一种有效技术。

---------- Backpropagation with Surrogate Gradients·END ----------


2.2 KNOWLEDGE DISTILLATION

知识蒸馏：利用教师模型的“基于响应”的知识（即软标签）来传递知识。

当这个概念第一次提出时，隐藏层中捕获的特征被忽略了，因为它们只关注当时的最终概率分布。为了更有效地学习教师模型，一些作品主张在蒸馏过程中纳入隐藏特征对齐。

此外，Park 等人引入了基于关系的知识蒸馏。证明训练数据示例之间的相互关系也很重要。

最近，有一些研究引入了知识蒸馏方法来训练 SNN。然而，它们大多数只关注图像分类任务，不能简单地应用于语言任务。在本研究中，我们提出了一种两阶段知识蒸馏方法来训练用于文本分类任务的 SpikeBERT，这是第一个展示将知识从大型语言模型转移到 SNN 的可行性的方法之一。


3.METHOD

在本节中，我们将描述如何改进 Spikformer 的架构，并介绍我们用于训练 SpikeBERT 的两阶段蒸馏方法。首先，我们将描述脉冲神经元和代理梯度在脉冲神经网络中的工作原理。然后我们将展示对 Spikformer 的简单但有效的修改，使其能够表示文本信息。最后，我们将详细说明“预训练特定任务”蒸馏。


3.1 SPIKING NEURONS AND SURROGATE GRADIENTS

漏积分激发（LIF）神经元（Wu et al., 2017）是最广泛使用的脉冲神经元之一。与 ReLU 等传统激活函数类似，LIF 神经元对输入的加权和进行操作，这有助于神经元在时间步tt的膜电位 UtU_t。如果神经元的膜电位达到阈值UthrU_{thr}，就会产生尖峰StS_t:

St={1if Ut≥Uthr0if Ut<UthrS_t= \begin{cases} 1& \text{if $U_t \geq U_{thr}$}\\ 0& \text{if $U_t < U_{thr}$} \end{cases}


我们可以将神经元膜电位的动态视为电阻电容电路（Maas，1997）。该电路微分方程的近似解可以表示为:

Ut=It+βUt−1−St−1Uthr,It=WXtU_t = I_t + \beta U_{t-1} - S_{t-1}U_{thr} , I_t = WX_t

  • XtX_t: LIF 神经元在时间步 t 的输入
  • W:一组用于整合不同输入的可学习权重，它是输入的加权和
  • β:膜电位的衰减率
  • Ut−1U_{t−1}:膜时间 t − 1 时的电位
  • St−1UthrS_{t−1}U_{thr}:模拟脉冲和膜电位重置机制

此外，我们关注 Fang 等人。 (2020a) 并使用类反正切代理梯度函数，该函数将 Heaviside 阶跃函数（方程 1）视为:

S≈1πarctan⁡(π2αU)+12S \approx \frac{1}{\pi}\arctan(\frac{\pi}{2}\alpha U) + \frac{1}{2}


上面等式的梯度为：

∂S∂U=α21(1+(π2αU)2)\frac{\partial S}{\partial U} = \frac{\alpha}{2} \frac{1}{(1+(\frac{\pi}{2}\alpha U)^2)}


3.2 THE ARCHITECTURE OF SPIKEBERT

SpikeBERT 是第一个用于语言任务的大型脉冲神经网络。我们的架构基于 Spikformer（Zhou 等人，2022），这是一种硬件友好的基于 Transformer 的脉冲神经网络，我们已在图 1 (a) 中展示了它。在 Spikformer 中，重要的模块是脉冲自注意力（SSA），它利用离散脉冲来近似普通的自注意力机制。它可以写成：

SSA(Qs,Ks,Vs)=S(BN(MLP(QsKsT∗Vs∗τ)))SSA(Q_s,K_s,V_s) = S(BN(MLP(Q_sK_s^T * V_s * \tau)))

Qs=SQs(BN(XsWQs)),Ks=SKs(BN(XsWKs)),Vs=SVs(BN(XsWVs))Q_s = S_{Q_s}(BN(X_sW_{Q_s})),K_s = S_{K_s}(BN(X_sW_{K_s})),V_s = S_{V_s}(BN(X_sW_{V_s}))

  • S: Heaviside 阶跃函数，如方程 1，Xs∈RT×L×DX_s ∈ R^{T×L×D} 是 SSA 的输入
  • T:时间步数
  • BN: 批量归一化
  • τ: 缩放因子
  • SSA和Qs、Ks、Vs的输出都是仅包含0和1的脉冲矩阵
  • WQs、WKs、WVsW_{Q_s}、W_{K_s}、W_{V_s}和MLP都是可学习的十进制参数
  • Spikformer 中注意力图的形状，即 QsKsTQ_sK_s^T ，是 D × D，其中 D 是隐藏层的维数

我们修改了 Spikformer，使其能够有效地处理文本数据。首先，我们用词嵌入层和脉冲神经元层替换脉冲补丁分割（SPS）模块，这是将标记映射到张量所必需的。最重要的是，我们认为通过注意力机制与不同位置的单词共享的特征比不同维度的单词共享的特征更重要。


Fig 2:用于训练 SpikeBERT 的两阶段蒸馏方法（预训练特定任务蒸馏）概述。

  • T是每层特征的时间步数
  • 仅在第 2 阶段考虑对数损失和交叉熵损失
  • 不同的颜色深浅表示浮点值的大小。
  • LfeaiL^i_{fea}下的虚线表示在计算特征对齐损失时可以忽略某些隐藏层的特征。如果学生模型包含的层数与教师模型不同，我们将每隔几层对齐特征。

因此，我们将脉冲自注意力（SSA）模块中的注意力图重塑为 N × N 而不是 D × D，其中 N 是输入句子的长度。最后，我们将“卷积层批量归一化”替换为“线性层归一化”，因为卷积层总是用于捕获图像中的像素特征，这不适合语言任务。我们在图 1 (b) 中展示了 SpikeBERT 的架构。


3.3 TWO-STAGE DISTILLATION

我们尝试过像 BERT 一样直接进行掩码语言建模（MLM）和下一句预测（NSP），但我们发现由于“自累积动态”，整个模型无法收敛（Fang et al., 2020b），这是大规模脉冲神经网络中尚未解决的问题。因此，我们选择使用知识蒸馏的方式来训练SpikeBERT，这样模型中代理梯度的偏差就不会快速累积。然而，将知识从 ANN 提炼到 SNN 是一个巨大的挑战，因为 ANN 和 SNN 中隐藏的特征不是相同的数据格式。 ANN 中的特征是浮点格式，而 SNN 中的特征是时变脉冲序列。我们发现在训练时引入外部模块（参见第3.3.1节）来对齐特征可以完美解决这个问题。

我们遵循流行的“训练前微调”配方，提出两阶段蒸馏。第一阶段是使用大规模语料库来对齐 BERT 和 SpikeBERT 之间的嵌入和隐藏特征。第二阶段是从经过微调的 BERT 到模型完成阶段 1 中提取任务特定数据集上的逻辑和交叉熵信息。我们在图 2 中展示了我们的方法的概述。


3.3.1 THE FIRST STAGE: PRE-TRAINING DISTILLATION

给定与下游任务无关的预训练 BERT（Devlin 等人，2019）作为教师 T M 和 SpikeBERT 作为学生 SM，我们此阶段的目标是将 TM 和 SM 的嵌入和隐藏特征与未标记的集合对齐文本。下面我们将介绍嵌入对齐损失和特征对齐损失:

特征对齐损失（Feature Alignment Loss）：这个损失 LfeaL_{fea} 是衡量每个隐藏层的 T M 和 SM 特征的相似度。然而，学生模型每层的特征 FsmF_{sm} 的形状为 T × N × D，而 BERT 的特征 FtmF_{tm} 的形状为 N × D，其中 T 是时间步数，D 是隐藏层的维数，L 是样本长度。更重要的是，FsmF_{sm} 是一个仅包含 0 和 1 的矩阵，而 FtmF_{tm} 是一个十进制矩阵。为了解决 FtmF_{tm} 和 FsmF_{sm} 之间的不同维度问题，以及 T M 的连续特征和 SM 的离散特征之间的差异问题，需要一种转换策略。我们遵循 Heo 等人的特征转换方法。 （2019）；陈等人。 （2021）；邱等人。 （2023）将TM和SM的特征映射到相同的内容空间：

Ftm′=Ftm,Fsm′=LayerNorm(MLP(∑tT(Fsmt)))F_{tm}^{'} = F_{tm},F^{'}_{sm} = LayerNorm(MLP(\sum_t^T(F^t_{sm})))

然而，我们发现在此阶段的前几层很难将学生模型生成的特征与 BERT 生成的特征对齐。我们认为这是因为学生模型可能需要更多的网络层来通过输入之间的交互来捕获基本特征。如图2所示，我们在计算特征对齐损失时选择忽略一些前面的层。假设BERT包含B Transformer块（即B层）并假设学生模型包含MSpike Transformer块。因此，如果 B > M ，我们将在每个⌈BM⌉\lceil \frac{B}{M} \rceil层对齐特征。对于学生模型中的第 i 层，其特征对齐损失为 Lfeai=∣∣Ftm′−Fsm′∣∣2L^i_{fea} = ||F^{'}_{tm} − F^{'}_{sm}||_2。

嵌入对齐损失（Embedding Alignment Loss） 正如 3.2 节中所讨论的，输入句子的嵌入在被前馈到 Heaviside 阶跃函数之前并不是脉冲的形式。将 $E_{tm} 和 EsmE_{sm} 分别定义为教师和学生的嵌入，因此嵌入对齐损失为

Lemb=∣∣Etm−MLP(Esm)∣∣2L_{emb} = ||E_{tm} − MLP(E_{sm})||_2

。 MLP 层是一个变换，其作用与公式 6 类似。

综上所述，在第 1 阶段，总损失 L1 是所选层的特征对齐损失之和：

L1=σ∑iLfeai+σ2LembL_1 = \sigma\sum_i L^i_{fea} + \sigma_2 L_{emb}

其中超参数 σ1 和 σ2 用于平衡嵌入和特征的学习。


3.3.2 THE SECOND STAGE: TASK-SPECIFIC DISTILLATION

在第 2 阶段，我们采用在特定任务数据集上微调的 BERT 作为教师模型，该模型作为学生完成第 1 阶段。为了完成某种语言任务，在基本语言模型上应该有一个特定于任务的头部，如图2所示。例如，需要在BERT上添加一个MLP层来进行文本分类。此外，数据增强是知识蒸馏中常用且高效的技术（Jiao等，2019；Tang等，2019；Liu等，2022）。下面，我们将讨论我们的数据增强方法，以及 logits 损失和交叉熵损失.

数据增强 在蒸馏方法中，小数据集可能不足以让教师模型充分表达其知识（Ba Caruana，2013）。为了解决这个问题，我们增加了训练集，以促进有效的知识蒸馏。我们跟随唐等人。 (2019) 增强训练集：首先，我们用概率pmaskp_{mask} 随机替换带有 [MASK] 标记的单词。其次，我们用相同 POS 标签的另一个单词替换一个单词，概率为 pposp_{pos}。第三，我们以概率 pngp_{ng} 从训练示例中随机采样 n-gram，其中 n 是从 {1, 2, …, 5} 中随机选择的。

Logits 损失 我们考虑了 logits,也称为软标签，它让学生学习老师的预测分布。为了衡量两个分布之间的距离，我们选择KL散度： Llogits=∑icpilog(piqi)L_{logits} = \sum^c_i p_ilog(\frac{p_i}{q_i}) ，其中c是类别数，pip_i和qiq_i表示教师模型和学生模型的预测分布。

交叉熵损失 交叉熵损失可以帮助学生模型从特定任务数据集中的样本中学习： Lce=−∑icq^ilog(qi)L_{ce} = − \sum^c_i \hat{q}_i log(qi)，其中 q^i\hat{q}_i 表示单热标签向量。因此，第 2 阶段的总损失 L2包含四个术语：

L2=λ1∑iLfeai+λ2Lemb+λ3Llogits+λ4LceL_2 = \lambda_1 \sum_i L^i_{fea} + \lambda_2 L_{emb} + \lambda_3 L_{logits} + \lambda_4 L_{ce}

  • λ1、λ2、λ3 和 λ4 是控制这些损失权重的超参数。
  • 对于这两个阶段，我们都采用时间反向传播（BPTT），它适合训练脉冲神经网络。