[論文レビュー] 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は、非微分可能なスパイクニューラルネットワーク（SNN）を回避するための暗黙の微分を用いて、逆誤差伝搬（BPTT）や補助勾配を用いずにエンドツーエンドの学習を可能にする、生物学的に妥当でエネルギー効率の良いスパイク言語モデルを提案する。新規のANN-SNN distillationフレームワークと、定常状態におけるスパイク発火レートに基づくスパイクアテンション機構を採用し、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.

研究の動機と目的

• 大規模言語モデル（LLM）のエネルギーコストを低減しつつ性能を維持する、生物学的に妥当で低消費電力のスパイク言語モデル（SNN）の開発。
• SNNの非微分性を補助勾配ではなく暗黙の微分を用いることで克服し、エンドツーエンドの学習を可能にする。
• 変換器ベースのアーキテクチャと互換性を持つために、定常状態における平均スパイク発火レート（ASR）を活用するスパイクアテンション機構の設計。
• 均衡状態におけるASRと中間層活性化パターンを用いて、事前学習済みのBERT「教師」モデルからスパイク「生徒」モデルへの効率的な知識蒸留を可能にする。
• GLUEベンチマーク上で動作する最初のスパイク言語モデルを実証し、エネルギー効率の面での実現可能性と性能向上を示す。

提案手法

• スパイクニューロンの定常状態収束から導かれる固定点方程式に暗黙の微分を適用し、時間に沿った逆誤差伝搬（BPTT）や補助勾配を回避する。
• クエリ、キー、バリューの演算をスパイク入力と均衡状態における平均スパイク発火レート（ASR）を用いて行うスパイクアテンション機構を導入し、標準アテンションと同等の性能を確保する。
• スパイク生徒モデルの中間層ASRを事前学習済みBERT「教師」モデルの活性化パターンと一致させる、新規のANN-SNN知識蒸留フレームワークを採用する。
• 段階的な蒸留を適用：一般蒸留、タスク固有の蒸留（例：MNLI、その後QNLI）、最終予測層の蒸留により、精度を向上させる。
• 正規化演算（Norm#OPS）と導出されたエネルギー効率要因 $e = ig(rac{1}{5.1} imes ext{Norm\#OPS}ig)^{-1}$ を用いてエネルギー削減を定量的に評価し、ACC演算はMACと比較してわずか0.9pJで、4.6pJに比べて5倍以上もエネルギー効率が良いことを示す。
• SST-2での電力消費、収束時間、精度のトレードオフを分析するため、$V_{ ext{th}}$ と $T_{ ext{conv}}$ におけるアブレーションスタディを実施する。

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を用いずに、スパイク変換器ベースの言語モデルを効果的に学習できるか？
• RQ2 スパイクベースでイベント駆動されるフレームワークで動作するスパイクアテンション機構をどのように設計すれば、標準自己アテンションの表現能力を維持できるか？
• RQ3 事前学習済みBERTモデルからの知識をASRベースの蒸留を用いてスパイク生徒モデルにどれほど効果的に転送できるか、そしてその結果、性能が向上するか？
• RQ4 SpikingBERTにおけるエネルギー-精度トレードオフはどのようなものか？非スパイクモデルと比較して、顕著なエネルギー効率の向上を達成できるか、かつ競争力ある精度を維持できるか？
• RQ5 ハイパーパrameterである$T_{ ext{conv}}$ と $V_{ ext{th}}$ が、スパイク言語モデルにおけるエネルギー消費、収束時間、モデル精度にどのように影響するか？

主な発見

• 知識蒸留を省略した場合、GLUEベンチマークで4–5%の精度低下が生じるが、これは本研究で提案する蒸留フレームワークの必要性を示している。
• T_{ ext{conv}} = 16 の場合、SpikingBERTはSST-2で最大精度の98%を達成し、同じサイズの非スパイクモデルと比較してほぼ2倍のエネルギー効率（$e$）を実現した。
• 膜電位しきい値（$V_{ ext{th}}$）を増加させることで、平均スパイク発火レート（ASR）と瞬時の電力消費が低下し、エッジデバイスに最適な低消費電力動作が可能になる。
• エネルギー効率要因 $e$ は $e = ig(rac{1}{5.1} imes ext{Norm\#OPS}ig)^{-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