Vision Transformer Acceleration via a Versatile Attention Optimization Framework

Vision Transformers (ViTs) have achieved remarkable success across various tasks. However, their deployment is hindered by challenges, such as high memory requirements, long inference latency, and significant power consumption. To address these challenges, existing works optimize one of the two key stages on ViT’s critical path: linear projection or self-attention. Regrettably, we have noticed that both linear projection and self-attention can potentially become bottlenecks as the input image resolution varies, which makes the existing approaches lack generality. Accordingly, in this article, we propose a versatile attention optimization framework. On the algorithm side, we present a SpQuant algorithm that sparsifies weight matrices offline and input matrices online during linear projection as well as tunes the bit-width of the probabilities matrix according to their importance. On the hardware side, we design SQArch architectures to improve the performance of the SpQuant algorithm. The proposed SQArch architecture offers a low-cost preprocess module that predicts and prunes nonkey elements of the input matrix on the fly. Moreover, we design a compute module that supports sparse-sparse matrix multiplications (SpMSpM) and multiple precision computations on a single systolic array for generality. Furthermore, we can address the underutilization and workload imbalance problems by 1) decoupling the rows in the systolic array for enough flexibility and 2) proposing a workload balance scheme for SpMSpM that allows the array to accept data of similar sparsity, thereby reducing synchronization between computing units. Extensive experiment results demonstrate that SQArch can achieve satisfactory performance speedups and energy saving compared to state-of-the-art designs.