Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Stochastic Spiking Attention: Accelerating Attention with Stochastic Computing in Spiking Networks

Zihang Song, Prabodh Katti, Osvaldo Simeone, Bipin Rajendran

TL;DR

A novel framework leveraging stochastic computing (SC) to effectively execute the dot-product attention for SNN-based Transformers and can achieve high classification accuracy on CIFAR-10 within 10 time steps, which is comparable to the performance of a baseline artificial neural network implementation.

Abstract

Spiking Neural Networks (SNNs) have been recently integrated into Transformer architectures due to their potential to reduce computational demands and to improve power efficiency. Yet, the implementation of the attention mechanism using spiking signals on general-purpose computing platforms remains inefficient. In this paper, we propose a novel framework leveraging stochastic computing (SC) to effectively execute the dot-product attention for SNN-based Transformers. We demonstrate that our approach can achieve high classification accuracy ($83.53\%$) on CIFAR-10 within 10 time steps, which is comparable to the performance of a baseline artificial neural network implementation ($83.66\%$). We estimate that the proposed SC approach can lead to over $6.3\times$ reduction in computing energy and $1.7\times$ reduction in memory access costs for a digital CMOS-based ASIC design. We experimentally validate our stochastic attention block design through an FPGA implementation, which is shown to achieve $48\times$ lower latency as compared to a GPU implementation, while consuming $15\times$ less power.

Stochastic Spiking Attention: Accelerating Attention with Stochastic Computing in Spiking Networks

TL;DR

A novel framework leveraging stochastic computing (SC) to effectively execute the dot-product attention for SNN-based Transformers and can achieve high classification accuracy on CIFAR-10 within 10 time steps, which is comparable to the performance of a baseline artificial neural network implementation.

Abstract

Spiking Neural Networks (SNNs) have been recently integrated into Transformer architectures due to their potential to reduce computational demands and to improve power efficiency. Yet, the implementation of the attention mechanism using spiking signals on general-purpose computing platforms remains inefficient. In this paper, we propose a novel framework leveraging stochastic computing (SC) to effectively execute the dot-product attention for SNN-based Transformers. We demonstrate that our approach can achieve high classification accuracy () on CIFAR-10 within 10 time steps, which is comparable to the performance of a baseline artificial neural network implementation (). We estimate that the proposed SC approach can lead to over reduction in computing energy and reduction in memory access costs for a digital CMOS-based ASIC design. We experimentally validate our stochastic attention block design through an FPGA implementation, which is shown to achieve lower latency as compared to a GPU implementation, while consuming less power.
Paper Structure (14 sections, 6 equations, 3 figures, 3 tables)

This paper contains 14 sections, 6 equations, 3 figures, 3 tables.

Table of Contents

  1. Introduction
  2. State of the Art
  3. Main Contributions
  4. Background
  5. Self-Attention
  6. Bernoulli Coding and Stochastic Computing
  7. Spiking Neurons
  8. Stochastic Spiking Attention
  9. Spike Coding of Inputs and Weight Matrices
  10. Stochastic Spiking Attention (SSA)
  11. Dataflow Design for SSA
  12. Optimizing for Hardware Implementation
  13. Experiments
  14. Conclusion

Figures (3)

  • Figure 1: Top: A conventional implementation of an attention block based on real-valued multiply-and-accumulate operations within an artificial neural network (ANN) architecture. Bottom: The proposed spiking neural network (SNN)-based attention block with spiking inputs, outputs, and stochastic computations. Multiplication operations are replaced with logical AND $(\land)$ operations on spikes. Further hardware efficiency is achieved by the replacement of scaling and softmax blocks with a Bernoulli rate encoder, as discussed in Section \ref{['sec:stoattn']}.
  • Figure 2: Top: The architectural schematic of the SSA block. Bottom: the $(i,j)$-th stochastic attention unit (SAU) illustrated in detail. All the wires, unless specified, carry one bit.
  • Figure 3: Illustration of dataflow design for the attention operation of each $(i,j)$-th stochastic attention unit (SAU).