[논문 리뷰] SpikingBERT: Distilling BERT to Train Spiking Language Models Using Implicit Differentiation

Malyaban Bal, Abhronil Sengupta|arXiv (Cornell University)|2023. 08. 21.

Advanced Memory and Neural ComputingEngineering인용 수 3

한 줄 요약

SpikingBERT는 스파iking 신경망의 비미분 가능성 문제를 회피하기 위해 암묵적 미분을 사용하여 순차적 미분을 피하고, 보조 기울기 없이 종단 간 훈련이 가능한 생물학적으로 타당한 에너지 효율적인 스파이킹 언어 모델을 제안한다. 새로운 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.

연구 동기 및 목표

• 대규모 언어 모델의 에너지 비용을 줄이면서도 성능을 유지하는 생물학적으로 타당한 저전력 스파이킹 언어 모델(SNN)을 개발하는 것.
• 스파이킹 신경망의 비미분 가능성 문제를 보완하기 위해 보조 기울기 대신 암묵적 미분을 사용하여 종단 간 훈련이 가능한 방법을 제안하는 것.
• 트랜스포머 기반 아키텍처와의 호환성을 확보하기 위해 정체 상태 평균 스파이크 빈도(ASR)를 활용하는 스파이킹 어텐션 메커니즘을 설계하는 것.
• 균형 상태 ASR과 중간 레이어 활성도를 기반으로 사전 훈련된 BERT '선생' 모델에서 스파이킹 '학생' 모델로 지식을 효율적으로 전달하는 방법을 제안하는 것.
• GLUE 벤치마크에서 작동 가능한 첫 번째 스파이킹 언어 모델을 구현하여 에너지 효율성 향상과 성능 향상을 입증하는 것.

제안 방법

• 스파이킹 뉴런의 정체 상태 수렴을 기반으로 한 고정점 방정식에 암묵적 미분을 적용하여, 시간에 따른 역전파(BPTT)와 보조 기울기를 피함.
• 질문, 키, 밸류 연산을 스파이크 기반 입력과 정체 상태에서의 평균 스파이크 빈도(ASR)를 사용하여 계산하는 스파이킹 어텐션 메커니즘을 도입하여 기존 어텐션과의 등가성을 확보함.
• 스파이킹 학생 모델의 중간 레이어 ASR를 사전 훈련된 BERT '선생' 모델의 활성도 패tern과 일치시키는 새로운 ANN-SNN 지식 전이 프레임워크를 도입함.
• 다단계 전이를 적용: 일반 전이, 과제별 전이(MNLI 이후 QNLI 등), 최종 예측 레이어 전이를 통해 정확도 향상.
• 정규화된 연산(Norm#OPS)과 유도된 에너지 효율성 요소 $e = ig(rac{1}{5.1} imes ext{Norm\#OPS}ig)^{-1}$ 를 사용하여 에너지 절감을 정량화함. ACC 연산은 0.9pJ로 MAC 대비 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 초기화 조건인 $T_{ ext{conv}}$와 $V_{ ext{th}}$는 스파이킹 언어 모델의 에너지 소모, 수렴 시간, 정확도에 어떤 영향을 미치는가?

주요 결과

• 지식 전이 프레임워크를 생략할 경우 GLUE 벤치마크에서 4–5%의 정확도 저하가 발생함으로써, 제안된 전이 프레임워크의 필요성을 입증함.
• T_{ ext{conv}} = 16일 때 SpikingBERT는 SST-2에서 최대 정확도의 98%를 달성하면서도 비스파이킹 모델 대비 거의 두 배의 에너지 효율성($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