[Paper Review] 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 ComputingEngineering3 citations

TL;DR

SpikingBERT proposes a bio-inspired, energy-efficient spiking language model trained via implicit differentiation to bypass the non-differentiability of spiking neural networks, enabling end-to-end training without surrogate gradients. It achieves competitive GLUE benchmark performance using a novel ANN-SNN distillation framework and a spiking attention mechanism based on steady-state spiking rates.

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.

Motivation & Objective

• To develop a biologically plausible, low-power spiking language model (SNN) that reduces the energy cost of large language models (LLMs) while maintaining performance.
• To overcome the non-differentiability of SNNs by using implicit differentiation instead of surrogate gradients, enabling end-to-end training.
• To design a spiking attention mechanism that leverages steady-state average spiking rates (ASR) for compatibility with transformer-based architectures.
• To enable efficient knowledge distillation from a pre-trained BERT 'teacher' to a spiking 'student' model using equilibrium ASR and intermediate layer activations.
• To demonstrate the first operational spiking LM on the GLUE benchmark, showing feasibility and performance gains in energy efficiency.

Proposed method

• Leverages implicit differentiation on fixed-point equations derived from the steady-state convergence of spiking neurons over $T_{ ext{conv}}$ time steps, avoiding backpropagation through time (BPTT) and surrogate gradients.
• Introduces a spiking attention mechanism where query, key, and value operations are computed using spike-based inputs and the average spiking rate (ASR) at equilibrium, ensuring equivalence to standard attention.
• Employs a novel ANN-SNN knowledge distillation framework that aligns intermediate layer ASR in the spiking student model with activation patterns in the pre-trained BERT 'teacher' model.
• Applies multi-stage distillation: general distillation, task-specific distillation (e.g., on MNLI then QNLI), and final prediction-layer distillation to improve accuracy.
• Uses normalized operations (Norm#OPS) and a derived energy-efficiency factor $e = ig(rac{1}{5.1} imes ext{Norm\#OPS}ig)^{-1}$ to quantify energy savings, with ACC operations consuming only 0.9pJ vs. 4.6pJ for MAC.
• Performs ablation studies on $V_{ ext{th}}$ and $T_{ ext{conv}}$ to analyze trade-offs between power consumption, convergence time, and accuracy on 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

Experimental results

Research questions

• RQ1 Can a spiking transformer-based language model be trained effectively without surrogate gradients or BPTT, using implicit differentiation on steady-state spiking dynamics?
• RQ2 How can a spiking attention mechanism be designed to preserve the representational capacity of standard self-attention while operating in a spike-based, event-driven framework?
• RQ3 To what extent can knowledge from a pre-trained BERT model be transferred to a spiking student model using ASR-based distillation, and does this improve performance?
• RQ4 What is the energy-accuracy trade-off in SpikingBERT, and can it achieve significant energy efficiency gains over non-spiking models while maintaining competitive accuracy?
• RQ5 How do hyperparameters like $T_{ ext{conv}}$ and $V_{ ext{th}}$ affect energy consumption, convergence time, and model accuracy in a spiking LM?

Key findings

• SpikingBERT achieves competitive performance on the GLUE benchmark, with a 4–5% accuracy drop when knowledge distillation is omitted, demonstrating the necessity of the proposed distillation framework.
• At $T_{ ext{conv}} = 16$, SpikingBERT attains 98% of the maximum accuracy on SST-2 while achieving nearly twice the energy efficiency ($e$) compared to a non-spiking model of the same size.
• Increasing the membrane potential threshold ($V_{ ext{th}}$) reduces average spiking rate (ASR) and instantaneous power consumption, enabling lower power operation ideal for edge devices.
• The energy-efficiency factor $e$ is derived as $e = ig(rac{1}{5.1} imes ext{Norm\#OPS}ig)^{-1}$, showing that ACC operations in SNNs are over 5× more energy-efficient than MAC operations.
• Task-specific knowledge distillation (e.g., MNLI → QNLI) improves performance over general distillation, with accuracy gains observed when distilling on downstream tasks.
• The spiking attention mechanism successfully emulates standard attention by aligning the steady-state ASR of spiking neurons with the attention output, enabling functional equivalence in the equilibrium state.

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