[論文レビュー] SpikeGPT: Generative Pre-trained Language Model with Spiking Neural Networks

Ruijie Zhu, Qihang Zhao+2 more|arXiv (Cornell University)|Feb 27, 2023

Advanced Memory and Neural Computing被引用数 52

ひとこと要約

SpikeGPTは、Spiking RWKVとSRFFNを組み合わせた、直接訓練された最初のスパイキングニューラルネットワーク(SNN)言語モデルで、線形時間計算量と neuromorphic のような疎活性化によるエネルギー削減を実現し、競争力のある生成と理解を達成します。

ABSTRACT

As the size of large language models continue to scale, so does the computational resources required to run it. Spiking Neural Networks (SNNs) have emerged as an energy-efficient approach to deep learning that leverage sparse and event-driven activations to reduce the computational overhead associated with model inference. While they have become competitive with non-spiking models on many computer vision tasks, SNNs have also proven to be more challenging to train. As a result, their performance lags behind modern deep learning, and we are yet to see the effectiveness of SNNs in language generation. In this paper, inspired by the Receptance Weighted Key Value (RWKV) language model, we successfully implement `SpikeGPT', a generative language model with binary, event-driven spiking activation units. We train the proposed model on two model variants: 45M and 216M parameters. To the best of our knowledge, SpikeGPT is the largest backpropagation-trained SNN model to date, rendering it suitable for both the generation and comprehension of natural language. We achieve this by modifying the transformer block to replace multi-head self attention to reduce quadratic computational complexity O(N^2) to linear complexity O(N) with increasing sequence length. Input tokens are instead streamed in sequentially to our attention mechanism (as with typical SNNs). Our preliminary experiments show that SpikeGPT remains competitive with non-spiking models on tested benchmarks, while maintaining 20x fewer operations when processed on neuromorphic hardware that can leverage sparse, event-driven activations. Our code implementation is available at https://github.com/ridgerchu/SpikeGPT.

研究の動機と目的

• Transformerに着想を得たアーキテクチャでスパイキングニューロンを直接用いて言語モデルをトレーニングすることを示す。
• 自己注意の二次計算量を線形に削減し、長いシーケンスとストリーミング入力を可能にする。
• ニューロモルフィックハードウェア上でのシナプス演算を抑えつつ、競争力のある自然言語生成と理解を達成する。
• 2つのパラメータスケール（45Mと216M）を示し、標準のTransformerベースのベースラインと比較する。
• NLPタスクにおける再帰とスパイキングダイナミクスを統合するトレーニング戦略とアーキテクチャ要素を探索する。

提案手法

• マルチヘッド自己注意を、線形時間計算を実現する再帰的に動作するSpiking RWKVトークンミキサーに置き換える。
• フィードフォワード成分として残差接続を持つSpiking RFFN チャンネルミキサーを使用する。
• 前方伝搬のためのHeavisideを用いた二値埋め込みとバックプロパゲーションのためのアーctangent代理勾配を導入する。
• 文脈を補完するためのトークンシフト操作を導入し、全注意は不要にする。
• Leaky Integrate-and-Fireニューロンを組み込み、二値スパイク出力を生成しストリーミング計算を可能にする。
• デコーダーのみで事前訓練し、その後NLGとNLUタスク向けにタスク固有の最上位層を適用してファインチューニングを行う。

実験結果

リサーチクエスチョン

• RQ1 スパイキングニューロン(SNN)を用いたTransformer風アーキテクチャで、生成的言語モデルを直接効果的に訓練できるか。
• RQ2 自己注意を線形・再帰的なSpiking RWKVに置換することで、競争力のある言語モデリングを実現しつつ計算コストを削減できるか。
• RQ3 SpikeGPTのニューロモルフィック様のスパースな活性化によるエネルギーと計算上の利点は、従来のTransformersと比べてどうか。
• RQ4 SpikeGPTの45Mおよび216Mパラメータ版は、標準的なNLG/NLUベンチマークでベースラインと比べてどうなるか。
• RQ5 SNNベースのNLPモデルに有効なトレーニング戦略（バイナリ埋め込み、代理勾配、トークンシフト）は何か。

主な発見

• SpikeGPTはこれまでで最大のバックプロパゲーション訓練済みSNN言語モデルであり（216Mパラメータ）、生成と理解のタスクで競争力のある性能を示す。
• SpikeGPTはSpiking RWKVを用いて二次式の注意計算を線形に削減し、ストリーミング的な逐次処理を行う。
• SpikeGPTはスパースでイベント駆動型の活性化により、標準のTransformerよりも SynOpsを20倍以上削減する。
• Enwik8では、SpikeGPT 45M（長さ1024）で訓練時/BPCが1.113/1.283、SpikeGPT 45M（長さ3072）で0.903/1.262を達成し、SynOpsはTransformerより顕著に低い。
• OpenWebText2で事前訓練した216MのSpikeGPTは、WikiText-2で競争的なパープレキシティを示し、WikiText-103ではGPT-2系と比較して相対的に低いパフォーマンスを示しており、スケールアップした訓練戦略の必要性を示唆する。
• SpikeGPTはエネルギー効率性を維持しつつ、競争力のあるNLGおよびNLUの結果を提供する。