ShiftAddViT: Mixture of Multiplication Primitives Towards Efficient Vision Transformer

Summary

This paper proposes an efficient accelerating framework for vision transformer, ShiftAddViT, which reparameterizes pre-trained ViTs with a mixture of complementary multiplication primitives and MoE designs. Specifically, All MatMuls in self-attention modules are reparameterized by additive kernels, and the remaining linear layers and MLPs are reparameterized by shift kernels. Besides, they develop a new MoE system for maintaining accuracy after reparameterization, and use a latency-aware load-balancing loss term to assign a dynamic amount of input tokens to each expert. Extensive experiments on various 2D/3D Transformer-based vision models demonstrate their superiority and efficiency.

Strengths

This paper is well-written and easy to follow. Converting multiplication operations to additive and shift is a promising technique in model compression and accleration, this paper provides a systematic design solutions for vision transformer architectures, which is interesting. Experiments in this paper are sufficient, including 2D and 3D task., and they provides detailed latency comparison for different methods, which much helps prove their effectiveness.

Weaknesses

ShiftAddViT is a complex system design solution, where the acceleration comes from the replacing of additive and shift operations with multiplication, the maintaining accuracy relies on the MoE design. And ShiftAddViT utils TVM to implement and optimize their customized kernels for practical hardware deployment on GPUs. In this paper, the contributions of multiplation less and MoE are much independent, this paper seems not much suitable for algorithm-based conferences like NeurIPS/CVPR/ICML.

Questions

1. What is the most important contributions in this paper？Multiplacation-less solution and MoE seems indepedent, It is a common sence that MoE could boost the performance on various tasks, improving accuracy with MoE seems unnecessary in this paper. 2. Re-parameterizing usually refers to one kind of mathematical equivalence transformations, such as RepVGG [1], AC-Net [2], what is the re-parameterizing in this paper? 3. ShiftAddNet [3] provides the first scheme for the replacing of shift and additive opration with multiplication, so what is the weakness for comination of ShiftAddNet and ViT architectures, could you show the comparision with ShiftAddViT. 4. When quantized to 4/8-bit, original ViT models could also achieve high performance and better acceleration, could you compare the quantized ViT with ShiftAddViT? 5. More complex task is better, such as detection and segementation. [1] Ding X, Zhang X, Ma N, et al. Repvgg: Making vgg-style convnets great again[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021: 13733-13742. [2] Ding X, Guo Y, Ding G, et al. Acnet: Strengthening the kernel skeletons for powerful cnn via asymmetric convolution blocks[C]//Proceedings of the IEEE/CVF international conference on computer vision. 2019: 1911-1920. [3]You H, Chen X, Zhang Y, et al. Shiftaddnet: A hardware-inspired deep network[J]. Advances in Neural Information Processing Systems, 2020, 33: 2771-2783.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

This paper provides the the limitation and societal impact discussion for their proposed technique.

Summary

The authors propose re-parameterizing ViTs to speed up inference without a full retraining. To do this, they introduce a new operation ShiftAddViT, which works well when applied to attention. For the MLP in a transformer, the authors use a mixture of experts. This operation reduces latency and energy usage, while maintaining accuracy.

Strengths

- The premise of leveraging power-of-two multiplications (a shift) is clever and a well-stated intuition (justified by Table 1 and also in L150-156). The additional energy analysis and chip-area statistics build ethos very well. - The paper presents a very thorough set of experiments on ImageNet classification and novel view synthesis. The latency analysis is thorough and well-documented / reasoned.

Weaknesses

- My only qualm is that (it appears) this relies on the existence of an effective binary quantization algorithm in L176. The results convince me that this binary quantization is acceptable, but this seems aggressive (as the authors mention later, when replacing the MLP). - From the introduction, I initially thought this re-parameterization would be arithmetically identical to the original. After reading the methods, it appears this isn't the case? In which case, wouldn't you need some fine-tuning? If so, how costly is the fine-tuning? If I'm wrong, this may be because I don't quite understand Figure 2 -- how is this related to the original matmul?

Questions

- I would've thought that the ShiftAddViT technique can really be applied to any operation (Clearly, that's not the case, as you added a modification for MLPs.) Why doesn't this work out of the box for any matrix multiplication? - What is "ideal parallelism" in the footnote for Table 4? Are these latencies not actually measured end-to-end? If this is the case, do the baseline latency measurements also use the same assumption of ideal parallelism? How fast/slow is the MoE kernel when plugged into the network and *entire network is measured end-to-end?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

See above. The evaluation is thorough, and the method is clearly presented. Despite the questions and objections I had above, I feel this is a cogently stated idea and clearly a thoroughly-investigated problem. I do still have some questions about the method and would like to know the answers, but the idea alone is definitely worth publishing.

Summary

This paper proposes a new type of multiplication-reduced model ShiftAddViT, which use the additive kernels to reparameterize the batched GEMM in the attention block and uses the shift kernels to reparameterize other MLPs or linear layers. In this way, it can reduce energy-intensive multiplications. The authors utilize TVM to implement and optimize those kernels and achieve a speedup on GPUs, and also energy savings.

Strengths

- The proposed ShiftAddViT achieves significant acceleration across different models on different tasks. - Experiments show the real speedup on GPUs to demonstrate the effectiveness of the proposed method.

Weaknesses

- Many implementation details of the proposed method are not clear enough. - It is unfair for the comparison with DeiT. The proposed method first modified the model architecture, and then introduced the shift and add kernel to improve the model efficiency. When compared with DeiT, the architectures of those two models are not the same. The comparison should be made under the condition that the model architectures are kept the same. - Based on Table 3 and Table 4, it appears that the “Shift” was not finally used in the classification task. And in Table 4, compared with the method that only enables “Add”, the final result (highlighted with red background) does not have any advantages in both model accuracy and latency. This seems to demonstrate that the “Shift” and “MoE” methods are not effective on classification tasks. - Table 4 and Table 5 lack the result of method “Quant.”+”Shift”+”MoE”. It would be better if these results could be provided to see the impact on model accuracy and latency when all the proposed methods are adopted.

Questions

About shift and add: - I think the Shift method and the power-of-two quantization should be similar. But there is no scaling factor in the Shift operation (equation 3). In this way, how to convert weights in pre-trained models to the shift layers format? Is the nearest neighbor method directly adopted? Wouldn't the reconstruction error be too large without a scaling factor? - How many bits of the P is used in the shift layers? And during finetuning, is the STE method used to solve the gradient backpropagation for s and P? - Both two binary methods KSH[32] and Quant[26] use the scaling factors, but those scaling factors are not shown in Figure 1. During inference, where to multiply those scaling factors? Just after the add operation? - How many bits are used for the input activations of the ShiftAddViT attention? Since only the fixed-point activations can be shifted, the input activation of shift layers needs to be quantized to an integer first. However, it seems that the article does not mention anything related to activation quantization. And there is no step related to the activation quantization in the two-stage finetuning process. In addition, if the activations are quantized, under the shift&add paradigm, where to multiply the scaling factor from activation quantization?  - How to quantize the input of the last Shift in ShiftAddViT Attention to integer values? Those input activations should also be fixed-point numbers, not floating-point. About the MoE framework: - In line 223, how to get the gate value $p_i$ based on the input token $x$? $p_i(x):= e^{p_i}/\Sigma_j^ne^{p_j}$ ? Is this a typo? About the experiments: - In Table 4, why the latency of PVT is larger than MHA? Is it because of the different batch sizes? For a fair comparison, it is better to use the same batch size here, since the model with batch_size=1 always gets worse throughput on GPUs. And it is better to show the speedup from the Linear Attention under the fair comparison. - What’s the difference between line 5 and line 10 in Table 4? They have the same configs, but different results. - Based on Table 4 and Table 6, “Shift” may cause an accuracy decrease, and it is necessary to introduce the “MoE” method to compensate for the model accuracy, which leads to an increase in latency. So overall, there is no advantage to using the “shift” method in terms of the model accuracy and latency. Does the “shift” method only for energy reduction? - Experiments use the Eyeriss accelerator to get energy consumption. What are the configs of this accelerator? What data type is used for this accelerator? When counting energy, is the energy of memory access counted?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The authors discussed the limitations of the implementation of their work. It is necessary to use a customized hardware accelerator to fully leverage the benefits of “Shift” and “Add”.

Summary

This paper introduces a novel reparameterization method for efficient Vision Transformers (ViT). The method replaces the heavy multiplication operations in ViT with a combination of shift and add operations. By mapping queries and keys to binary codes in Hamming space and reparameterizing multi-layer perceptrons (MLPs) or linear layers, the multiplication operations in the model are effectively reduced. Experimental results demonstrate that ShiftAddViT achieves efficient performance on various 2D/3D Transformer visual tasks while achieving latency reduction and energy savings.

Strengths

This paper presents a clear motivation for each proposed compound. For example, the reparameterization of multiplications is inspired by hardware design practices and the concept of multiple experts (MOE) for dynamically routing tokens into multiplication or shift groups. The three questions posed by the authors clearly indicate the specific problems they are addressing. The writing and demonstration are clear and concise. In terms of results, this work delivers some impressive findings in regard to both efficiency and accuracy. Extensive ablation experiments are conducted to validate the proposed approach further.

Weaknesses

1. A primary concern is the limited validation of the method proposed by the authors, which is solely performed on small-scale models, with the largest model encompassing 30M and 4G FLOPs. I am particularly interested in the performance of larger-scale models. Naturally, I acknowledge the constraints imposed by computational resources. However, considering the hardware specifications disclosed by the authors, it is apparent that training models of ViT-Base size or similar are feasible. 2. Secondly, my concern lies in the stability of the training process and the associated wall-clock time. Obtaining information regarding the overall training cost and wall-clock time from the authors would offer valuable insights into the algorithm's feasibility and cost implications.

Questions

Overall, I am satisfied with the results and motivation. I am particularly concerned about whether this method can improve inference performance for larger models. The bottleneck for small models is not particularly high, to begin with, and compared to larger models, the performance of small models will undoubtedly be lower. Therefore, if the authors' method can enable larger models to achieve a significant increase in inference speed, it would greatly enhance the significance of the approach.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

The author discussed the limitation in their paper.

Summary

This work proposes ShiftAddViT, which is an efficient ViT reparameterization with a mixture of complementary multiplication primitives, such as bitwise shifts and adds. The alternative parameterization (quantization) is carefully examined to allocate to different components (MHSA, MLP) in ViT. In addition, the authors propose a mixture of experts (MoE) framework to classify input tokens and assign different primitives to best preserve accuracy. The MoE framework is guided by latency-aware load-balancing loss. Multiple empirical results are demonstrated including 2D ViT in image classification task, as well as GNT for NVS task.

Strengths

1. The paper is clearly written and easy to follow. 2. Though shift and add reparameterization is not completely new (binary, ternary or power-of-2 quantization), this is a very pioneer work to systematically investigate layer sensitivity and accuracy impact, allocation strategy, and hardware implementations. The hardware benchmarks with considerable latency and energy savings are impressive. 3. The experimental analysis is strong. In addition to conventional classification task, the authors also provide results on NVS task.

Weaknesses

1. Since this work proposes a collection of analysis and optimizations on MHSA and MLP, reducing the computation complexity and energy consumption, I wonder if this work is portable to NLP tasks, especially for LLMs with more tokens? 2. I cannot fully follow the token dispatching method mentioned in Section 4.2 and Figure 6. The authors did not discuss the methodology for input allocation and parallel computing issues in detail (line#227). Intuitively, if the important and sensitive tokens (yellow) in Figure 6 are fed into powerful experts (MULT.), while the rest are into SHIFT experts, are they still visible to each other (i.e. global receptive field)?

Questions

Please refer to concerns raised in weaknesses.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

N/A

Thanks for the response

Thanks to the authors for their thorough response; this was very helpful, and I hadn't realized the MoE was a mixture of multiplication *or add experts. The fine-tuning cost analysis was also welcome, although I also recognize that inference is the focus here. The above addresses my questions and concerns, so I stand by my rating! (I think you can just submit multiple "rebuttals" per review, so the length limit per response doesn't really matter.)

Response to rebuttal

Thanks for your response. The rebuttal is clear and easy to understand, all my concerns have been addressed. This work is good enough to be published.

Thanks for the rebuttal

Thanks for the larger models' experiments. Although we can see the performance has a potential drop, the improvement in inference time is notable. My concerns have been mainly addressed. I have raised my score.