SlowGuess/NeurIPS_2023_OCR huggingface dataset:
ShiftAddViT: Mixture of Multiplication Primitives Towards Efficient Vision Transformers
Haoran You*, Huihong Shi*, Yipin Guo*, and Yingyan (Celine) Lin
Georgia Institute of Technology, Atlanta, GA
{haoran.you, celine.lin}@gatech.edu, eic-lab@groups.gatech.edu
Vision Transformers (ViTs) have shown impressive performance and have become a unified backbone for multiple vision tasks. However, both the attention mechanism and multi-layer perceptrons (MLPs) in ViTs are not sufficiently efficient due to dense multiplications, leading to costly training and inference. To this end, we propose to reparameterize pre-trained ViTs with a mixture of multiplication primitives, e.g., bitwise shifts and additions, towards a new type of multiplication-reduced model, dubbed ShiftAddViT, which aims to achieve end-to-end inference speedups on GPUs without requiring training from scratch. Specifically, all MatMuls among queries, keys, and values are reparameterized using additive kernels, after mapping queries and keys to binary codes in Hamming space. The remaining MLPs or linear layers are then reparameterized with shift kernels. We utilize TVM to implement and optimize those customized kernels for practical hardware deployment on GPUs. We find that such a reparameterization on (quadratic or linear) attention maintains model accuracy, while inevitably leading to accuracy drops when being applied to MLPs. To marry the best of both worlds, we further propose a new mixture of experts (MoE) framework to reparameterize MLPs by taking multiplication or its primitives as experts, e.g., multiplication and shift, and designing a new latency-aware load-balancing loss. Such a loss helps to train a generic router for assigning a dynamic amount of input tokens to different experts according to their latency. In principle, the faster the experts run, the more input tokens they are assigned. Extensive experiments on various 2D/3D Transformer-based vision tasks consistently validate the effectiveness of our proposed ShiftAddViT, achieving up to  $5.18 \times$  latency reductions on GPUs and  $42.9\%$  energy savings, while maintaining a comparable accuracy as original or efficient ViTs. Codes and models are available at https://github.com/GATECH-EIC/ShiftAddViT.
Vision Transformers (ViTs) have emerged as powerful vision backbone substitutes to convolutional neural networks (CNNs) due to their impressive performance [57, 16]. However, ViTs' state-of-the-art (SOTA) accuracy comes at the price of prohibitive hardware latency and energy consumption for both training and inference, limiting their prevailing deployment and applications [34, 51], even though pre-trained ViT models are publicly available. The bottlenecks are twofold based on many profiling results: attention and MLPs, e.g., they account for  $36\% / 63\%$  of the FLOPs and  $45\% / 46\%$  of the latency when running DeiT-Base [55] on an NVIDIA Jetson TX2 GPU [72, 19]). Previous efficient ViT solutions mostly focus on macro-architecture designs [21, 39, 32, 55, 9, 61] and linear attention optimization [59, 50, 5, 13, 6, 73]. However, they pay less attention on reducing the underlying dominant multiplications as well as on co-optimizing both attention and MLPs.
We identify that one of the most effective yet still missing opportunities for accelerating ViTs is to reparameterize their redundant multiplications with a mixture of multiplication primitives, i.e.,
bitwise shifts and adds. The idea is drawn from the common hardware design practice in computer architecture or digital signal processing. That is, the multiplication can be replaced by bitwise shifts and adds [67, 22]. Such a hardware-inspired approach can lead to an efficient and fast hardware implementation without compromising the model accuracy. As such, this paper embodies the shift&add idea in ViTs towards a new type of multiplication-reduced model. Moreover, unlike previous ShiftAddNet [70] which requires full and slow training and dedicated hardware accelerator support, this work starts from pre-trained ViTs to avoid the tremendous cost of training from scratch and targets end-to-end inference speedups, i.e., accelerating both attention and MLPs, on GPUs.
The aforementioned shift&add concept uniquely inspires our design of naturally hardware-friendly ViTs, but it also presents us with three challenges to address. First, how to effectively reparameterize ViTs with shifts and adds? Previous ShiftAddNet [70] reparameterizes CNNs by cascading shift layers and add layers, leading to a doubled number of layers or parameters. The customized CUDA kernels of shift and add layers also suffer from much slower training and inference than PyTorch [42] on GPUs. Both motivate a new easy-to-accelerate reparameterization method for ViTs. Second, how to maintain the accuracy after reparameterization? It is expected to see accuracy drops when multiplic
tions are turned into shift and add primitives [8, 11]. Most works compensate for the accuracy drops by enlarging model sizes leveraging the much improved energy efficiency advantages of shifts and adds [70], e.g., up to  $196\times$  unit energy savings than multiplications as shown in Tab. 1. While in ViTs, where input images are split into non-overlapping tokens, one can uniquely leverage the inherent adaptive sensitivity among input tokens. In principle, the essential tokens containing target objects are expected to be processed using more powerful multiplications. In contrast, tokens with less important background information can be handled by cheaper primitives. Such a principle aligns with the spirit of recent token merging [3] and input adaptive methods [45, 68, 74] for ViTs. Third, how can we balance the loading and processing time for the aforementioned important and less-important input tokens? For ViTs, using a mixture of multiplication primitives can lead to varying processing speeds for sensitive and non-sensitive tokens. This needs to be balanced; otherwise, intermediate activations may take longer to synchronize before progressing to the next layer.
Table 1: Hardware cost under  ${45}\mathrm{\;{nm}}$  CMOS.
OPs	Format	Energy (pJ)	Area (μm2)
Mult.	FP32	3.7	7700
FP16	0.9	1640
INT32	3.1	3495
INT8	0.2	282
Add	FP32	1.1	4184
FP16	0.4	1360
INT32	0.1	137
INT8	0.03	36
Shift	INT32	0.13	157
INT16	0.057	73
INT8	0.024	34
To the best of our knowledge, this is the first attempt to address the above three challenges and to incorporate the shift&add concept from the hardware domain in designing multiplication-reduced ViTs using a mixture of multiplication primitives. Specifically, we make the following contributions:
- We take inspiration from hardware practice to reparameterize pre-trained ViTs with a mixture of complementary multiplication primitives, i.e., bitwise shifts and adds, to deliver a new type of multiplication-reduced network, dubbed ShiftAddViT. All MatMuls in attention are reparameterized by additive kernels, and the remaining linear layers and MLPs are reparameterized by shift kernels. The kernels are built in TVM for practical deployments.
- We introduce a new mixture of experts (MoE) framework for ShiftAddViT to preserve accuracy post-reparameterization. Each expert represents either a multiplication or its primitives, such as shifts. Depending on the importance of a given input token, the appropriate expert is activated, for instance, multiplications for vital tokens and shifts for less crucial ones.
- We introduce a latency-aware load-balancing loss term within our MoE framework to dynamically allocate input tokens to each expert. This ensures that the number of tokens assigned aligns with the processing speeds of the experts, significantly reducing synchronization time.
- We conduct extensive experiments to validate the effectiveness of our proposed ShiftAddViT. Results on various 2D/3D Transformer-based vision models consistently show its superior efficiency, achieving up to  $5.18 \times$  latency reductions on GPUs and  $42.9\%$  energy savings, while maintaining a comparable or even higher accuracy than the original ViTs.
Vision Transformers (ViTs). ViTs [16, 55] split images into non-overlapping patches or tokens and have outperformed CNNs on a range of vision tasks. They use a simple encoder-only architecture, primarily made up of attention modules and MLPs. The initial ViTs necessitated costly pre-training
on expansive datasets [52], and subsequent advancements like DeiT [55] and T2T-ViT [75] alleviated this by refining training methodologies or crafting innovative tokenization strategies. Later on, a surge of new ViT architecture designs, such as PVT [60], CrossViT [7], PiT [26], MViT [18], and SwinTransformer [35], has emerged to improve the accuracy-efficiency trade-offs using a pyramid-like architecture. Another emergent trend revolves around dynamic ViTs like DynamicViT [45], A-ViT [68], ToME [3], and MIA-Former [74], aiming to adaptively eliminate non-essential tokens. Our MoE framework resonates with this concept but addresses those tokens using low-cost multiplication primitives that run concurrently. For edge deployment, solutions such as LeViT [21], MobileViT [39], and EfficientViT [5, 50] employ efficient attention mechanisms or integrate CNN feature extractors. This approach propels ViTs to achieve speeds comparable to CNNs, e.g., MobileNets [47]. In contrast, our proposed ShiftAddViT is the first to draw from shift&add hardware shortcuts for reparameterizing pre-trained ViTs, boosting end-to-end inference speeds and energy efficiency without hurting accuracy. Therefore, ShiftAddViT is orthogonal to existing ViT designs and can be applied on top of them.
Tremendous efforts have been dedicated to designing efficient ViTs. For instance, to address the costly self-attention module, which possesses quadratic complexity in relation to the number of input tokens, various linear attentions have been introduced. These can be broadly divided into two categories: local attention [35, 2, 56] and kernel-based linear attention [30, 13, 65, 38, 5, 34, 1]. For the former category, Swin [35] computes similarity measures among neighboring tokens rather than all tokens. QnA [2] disseminates the attention queries across all tokens. MaxViT [56] employs block attention and integrates dilated global attention to capture both local and global information. Regarding the latter category, most methods approximate the softmax function [13, 30, 4] or the complete self-attention matrix [65, 38] using orthogonal features or kernel embeddings, allowing the computational order to switch from (QK)V to Q(KV). Furthermore, a few efforts target reducing the multiplication count in ViTs. For instance, Ecoformer [34] introduces a novel binarization paradigm through kernelized hashing for Q/K, ensuring that MatMuls in attentions become solely accumulations. Adder Attention [51] examines the challenges of integrating adder operations into attentions and suggests including an auxiliary identity mapping. ShiftViT [58] incorporates spatial shift operations into attention modules. Different from previous works, our ShiftAddViT focuses on not only attention but also MLPs towards end-to-end ViT speedups. All MatMuls can be transformed into additions/accumulations using quantization, eliminating the need for complex kernelized hashing (KSH) as in [34]. For the remaining linear layers in the attention modules or MLPs, they are converted to bitwise shifts or MoEs, rather than employing spatial shifts like [58] for attentions.
Multiplication-less NNs. Numerous efforts aim to minimize the dense multiplications, which are primary contributors to the time and energy consumption in CNNs and ViTs. In the realm of CNNs, binary networks [14, 29] binarize both activations and weights; AdderNet [8, 66, 62] fully replaces multiplications with additions with a commendable accuracy drop; Shift-based networks use either spatial shifts [64] or bitwise shifts [17] to replace multiplications. Recently, techniques initially developed for CNNs have found applications in ViTs. For instance, BiTs [36, 25] introduce binary transformers that maintain commendable accuracy; Shu et al. and Wang et al. also migrate the add or shift idea to ViTs with the scope limited to attentions [51, 58]. Apart from network designs, there are also neural architecture search efforts made towards both accurate and efficient networks [71, 31]. As compared to the most related work, i.e., ShiftAddNet [70] and ShiftAddNAS [71], our ShiftAddViT for the first time enables shift&add idea in ViTs without compromising the model accuracy, featuring three more characteristics that make it more suitable for practical applications and deployment: (1) starting from pre-trained ViTs to avoid the costly and tedious training from scratch; (2) focusing on optimization for GPUs rather than on dedicated hardware such as FPGAs/ASICs; (3) seamlessly compatible with the MoE framework to enable switching between the mixture of multiplication primitives, such token-level routing and parallelism are uniquely applicable and designed for ViTs.
Self-attention and Vision Transformers. Self-attention is a core ingredient of Transformers [57, 16], and usually contains multiple heads H with each measuring pairwise correlations among all input tokens to capture global-context information. It can be defined as below:
$$
\operatorname {A t t n} (\mathbf {X}) = \operatorname {C o n c a t} \left(\mathrm {H} _ {1}, \dots , \mathrm {H} _ {h}\right) \mathbf {W} _ {O}, \text {w h e r e} \mathbf {H} _ {i} = \operatorname {S o f t m a x} \left(\frac {f _ {Q} (\mathbf {X}) \cdot f _ {K} (\mathbf {X}) ^ {T}}{\sqrt {d _ {k}}}\right) \cdot f _ {V} (\mathbf {X}), \tag {1}
$$
where  $h$  denotes the number of heads. Within each head, input tokens  $\mathbf{X} \in \mathbb{R}^{n \times d}$  of length  $n$  and dimension  $d$  will be linearly projected to query, key, and value matrices, i.e.,  $\mathbf{Q}, \mathbf{K}, \mathbf{V} \in \mathbb{R}^{n \times d_k}$ ,
Figure 1: An illustration of the network architecture overview of ShiftAddViT.
through three linear mapping functions,  $f_{Q} = \mathbf{X}\mathbf{W}_{Q}, f_{K} = \mathbf{X}\mathbf{W}_{K}, f_{V} = \mathbf{X}\mathbf{W}_{V}$ , where  $d_{k} = d / h$  is the embedding dimension of each head. The results from all heads are concatenated and projected with a weight matrix  $\mathbf{W}_O \in \mathbb{R}^{d \times d}$  to generate the final outputs. Such an attention module is followed by MLPs with residuals to construct one transformer block that can be formulated as below:
$$
\mathbf {X} _ {\text {A t t n}} = \operatorname {A t t n} (\operatorname {L a y e r N o r m} (\mathbf {X})) + \mathbf {X}, \quad \mathbf {X} _ {\text {M L P}} = \operatorname {M L P} (\operatorname {L a y e r N o r m} (\mathbf {X} _ {\text {A t t n}})) + \mathbf {X} _ {\text {A t t n}}. \tag {2}
$$
Shift and Add Primitives. Direct hardware implementation of multiplications is inefficient. Shift and add primitives serve as "shortcuts" for streamlined hardware designs. Shift operations, equivalent to multiplying by powers of two, offer significant savings. As shown in Tab. 1, such shifts can reduce energy use by up to  $23.8 \times$  and chip area by  $22.3 \times$  compared to multiplications [33, 70], given the INT32 data format and  $45\mathrm{nm}$  CMOS technology. Similarly, add operations, another efficient primitive, can achieve up to  $31.0 \times$  energy and  $25.5 \times$  area savings relative to multiplications. Both primitives have driven numerous network architecture innovations [8, 17, 70, 71].
4 The Proposed ShiftAddViT
Overview. Given the widespread availability of pre-trained ViT model weights, we aim to fine-tune and transform them into multiplication-reduced ShiftAddViTs in order to decrease runtime latency and enhance energy efficiency. To achieve this, we must reparameterize both the attentions and MLPs in ViTs to shift and add operations. Instead of using cascaded shift layers and add layers as in ShiftAddNet [70] or designing specific forward and backpropagation schemes like AdderNet [8, 51], we lean on the simple yet effective approach of ViTs' long-dependency modeling capability. We propose replacing the dominant MLPs/Linears/MatMu1s with shifts and adds, leaving the attention intact. This approach permits ShiftAddViTs to be adapted from existing ViTs without the need for training from scratch. As illustrated in Fig. 1, for attentions, we transform four Linear layers and two MatMu1s into shift and add layers, respectively. For MLPs, direct replacement with shift layers results in significant accuracy degradation. Hence, we consider designing a dedicated MoE framework incorporating a blend of multiplication primitives, such as multiplication and shift, to achieve both accurate and efficient final models. These steps transition pre-trained ViTs into ShiftAddViTs, achieving much reduced runtime latency while maintaining a comparable or even higher accuracy.
4.1 ShiftAddViT: Reparameterization Towards Multiplication-reduced Networks
This subsection describes how we reparameterize pre-trained ViTs with hardware-efficient shift and add operations, including a detailed implementation and the corresponding sensitivity analysis.
Reparameterization of Attention. There are two categories of layers in attention modules: MatMuls and Linears, that can be converted to Add and Shift layers, respectively. In order to reparameterize MatMuls to Add layers, we consider performing binary quantization on one operand during MatMuls, e.g.,  $\mathbf{Q}$  or  $\mathbf{K}$  for the MatMul of QK, so that the multiply-and-accumulate (MAC) operations between two matrices will be replaced with merely energy-efficient add operations [34]. Furthermore, to build ShiftAddViT on top of efficient linear attentions [59, 34, 38], we exchange the order of MatMuls from (QK)V to  $\mathbf{Q}(\mathbf{KV})$  for achieving linear complexity w.r.t. the number of input tokens. In this way, the binary quantization will be applied to  $\mathbf{K}$  and  $\mathbf{Q}$  as illustrated in Fig. 1 (b), leaving the more sensitive V branch as high precision. This also allows us to insert a lightweight depthwise convolution (DWConv) to  $\mathbf{V}$  branch in parallel to enhance local feature capturing capability with negligible
Figure 2: Illustration of both Shift and Add layers, where  $\mathbf{X}$  and  $\mathbf{O}$  refer to inputs and outputs.
overhead ( $< 1\%$  of total MACs) following recent SOTA linear attention designs [65, 73]. On the other hand, for reparameterizing the left four Linears in the attention module with Shift layers, we resort to sign flips and power-of-two shift parameters to represent the Linear weights [17, 70].
We illustrate the implementation of Shift and Add layers in Fig. 2 and formulate them as follows:
$$
\mathbf {O} _ {\text {S h i f t}} = \sum \mathbf {X} ^ {\mathrm {T}} \cdot \mathbf {W} _ {\mathrm {S}} = \sum \mathbf {X} ^ {\mathrm {T}} \cdot \mathbf {s} \cdot 2 ^ {\mathrm {P}}, \quad \mathbf {O} _ {\text {A d d}} = \sum \mathbf {X} ^ {\mathrm {T}} \cdot \mathbb {1} \left\{\mathbf {W} _ {\mathbf {A}} \neq \mathbf {0} \right\}, \tag {3}
$$
where  $\mathbf{X}$  refers to input activations,  $\mathbf{W}_{\mathbf{S}}$  and  $\mathbf{W}_{\mathbf{A}}$  are weights in Shift and Add layers, respectively. For Shift weights  $\mathbf{W}_{\mathbf{S}} = \mathbf{s} \cdot 2^{\mathbf{P}}$ ,  $\mathbf{s} = \text{sign}(\mathbf{W}) \in \{-1,1\}$  denotes sign flips,  $\mathbf{P} = \text{round}(\log_2(\text{abs}(\mathbf{W})))$ , and  $2^{\mathbf{P}}$  represents bitwise shift towards left ( $\mathbf{P} > 0$ ) or right ( $\mathbf{P} < 0$ ). Both  $\mathbf{s}$  and  $\mathbf{P}$  are trainable during the training. For Add weights, we have applied binary quantization and will skip zeros directly, the remaining weights correspond to additions or accumulations. Combining those two weak players, i.e., Shift and Add, leads to higher representation capability as also stated in ShiftAddNet [70], we highlight that our ShiftAddViT seamlessly achieves such combination for ViTs, without the necessity of cascading both of them as a substitute for one convolutional layer.
Reparameterization of MLPs. MLPs are often overlooked but also dominate ViTs' runtime latency apart from attentions, especially for the small or medium amount of tokens [19, 78, 43]. One natural thought is to reparameterize all MLPs with Shift layers as defined above. However, such aggressive design indispensably leads to severe accuracy drops, e.g.,  $\downarrow 1.18\%$  for PVTv1-T [60], (DWConv used between two MLPs in PVTv2 are kept). We resolve this by proposing a new MoE framework to achieve speedups with comparable accuracy, as shown in Fig. 1 (c) and will be elaborated in Sec. 4.2.
Sensitivity Analysis. To investigate whether the above reparameterization methods work as expected, we break down all components and conduct sensitivity analysis. Each of them is applied to ViTs with only 30 epochs of finetuning. As summarized in Tab. 2, applying linear attention (LA), Add, or Shift to attention layers makes little influence on accuracy. But applying Shift to MLPs may lead to severe accuracy drops for sensitive models like PVTv1-T. Also, making MLPs more effi
Table 2: Sensitivity analysis of reparameterizing ViTs with Shift and Add.
Components	Apply	Acc. (%) of PVT [60, 61]
PVTv2-B0	PVTv1-T
-	-	70.50	75.10
-	MSA	71.25	76.21
Attention	LA+Add	70.95	75.20
Shift	70.96	76.05
MLPs	Shift	70.28	73.92
MoE	70.86	74.81
cient contribute a lot to the energy efficiency (e.g., occupy  $65.7\%$  on PVTv2-B0). Both the accuracy and efficiency perspectives motivate us to propose a new reparameterization scheme for MLPs.
4.2 ShiftAddViT: Mixture of Experts Framework
This subsection elaborates on the hypothesis and method of our ShiftAddViT MoE framework.
Hypothesis. We hypothesize there are important yet sensitive input tokens that necessitate the more powerful networks otherwise suffer from accuracy drops. Those tokens are likely to contain object information that directly correlates with our task objectives as validated by [3, 74, 45, 68]. While the left tokens are less sensitive and can be fully represented even using cheaper Shift layers. Such an input-adaptive nature of ViTs calls for a framework featuring a mixture of multiplication primitives.
Mixture of Multiplication Primitives. Motivated by the sensitivity analysis in Sec. 4.1, we consider two experts Mult. and Shift for processing important yet sensitive tokens and insensitive tokens, respectively. As illustrated in Fig. 1 (c), we apply the MoE framework to compensate for the accuracy drop. Each input token representation  $\mathbf{x}$  will be passed to one expert according to the gate values  $p_i = G(\mathbf{x})$  in routers, which are normalized via a softmax distribution  $p_i \coloneqq \frac{e^{p_i}}{\sum_j^n e^{p_j}}$  and will be jointly trained with model weights. The output can be formulated as  $\mathbf{y} = \sum_{i}^{n} G(\mathbf{x}) \cdot E_i(\mathbf{x})$ , where gate function  $G(\mathbf{x}) = p_i(\mathbf{x}) \cdot \mathbb{1}\{p_i(\mathbf{x}) \geq p_j(\mathbf{x}), \forall j \neq i\}$ ,  $n$  and  $E_i$  denote the number of experts and  $i$ -th expert, respectively. Specifically, there are two obstacles when implementing such MoE in TVM: