ShiftAddViT: Mixture of Multiplication Primitives Towards Efficient Vision Transformer

Overview

• 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.
• The authors propose a new model called ShiftAddViT that reparameterizes pre-trained ViTs with a mixture of multiplication primitives, such as bitwise shifts and additions, to achieve end-to-end inference speedups on GPUs without requiring training from scratch.

Plain English Explanation

Vision Transformers (ViTs) are a type of deep learning model that have shown great success in a variety of visual tasks. However, the key components of ViTs, the attention mechanism and multi-layer perceptrons (MLPs), are computationally expensive due to the large number of multiplications involved. This makes ViTs costly to train and run efficiently on hardware like GPUs.

The researchers propose a new model called ShiftAddViT that aims to address this issue. Instead of using standard multiplications, ShiftAddViT reparameterizes the attention mechanism to use simpler operations like bitwise shifts and additions. This reduces the overall computational burden without significantly impacting the model's accuracy.

For the remaining MLP layers, the researchers develop a new "mixture of experts" framework that dynamically assigns different parts of the input to specialized sub-modules that use either multiplication or shift-based operations, depending on which is more efficient. This helps maintain the model's performance while further accelerating inference.

The key idea is to transform the original ViT model into a more efficient version, ShiftAddViT, that can run much faster on GPUs without having to retrain the entire model from scratch. This could lead to significant improvements in the real-world deployment of ViTs.

Technical Explanation

The authors propose a new model, dubbed ShiftAddViT, that reparameterizes pre-trained Vision Transformers (ViTs) to replace the computationally expensive matrix multiplications (MatMuls) with more efficient operations like bitwise shifts and additions.

Specifically, the attention mechanism in ViTs is reparameterized by mapping the queries and keys to binary codes in Hamming space, allowing the MatMuls among queries, keys, and values to be replaced with additive kernels. The remaining MLP or linear layers are then reparameterized with shift kernels.

To handle the accuracy drops caused by the MLP reparameterization, the authors propose a new "mixture of experts" (MoE) framework. This framework dynamically assigns input tokens to different expert sub-modules that use either multiplication or shift-based operations, depending on their latency. A new latency-aware load-balancing loss function is introduced to train the routing mechanism.

The authors utilize the TVM framework to implement and optimize these customized kernels for practical hardware deployment on GPUs. Extensive experiments on various 2D and 3D Transformer-based vision tasks demonstrate that ShiftAddViT can achieve up to 5.18x latency reductions on GPUs and 42.9% energy savings, while maintaining comparable accuracy to the original or efficient ViT models.

Critical Analysis

The paper presents a clever approach to reparameterizing ViT models to reduce the computational burden of the attention mechanism and MLPs, which are the main bottlenecks in these models. The proposed ShiftAddViT model and the mixture of experts framework for handling MLP reparameterization are both interesting and well-designed solutions.

However, the authors acknowledge that the MLP reparameterization does lead to some accuracy drops, which they try to mitigate with the MoE approach. It would be valuable to see more analysis on the trade-offs between the achieved inference speedups and the impact on model performance across different tasks and datasets.

Additionally, the paper does not provide much insight into the generalizability of the proposed techniques. It would be helpful to understand how well the ShiftAddViT approach might work for other types of Transformer-based models beyond just ViTs, or if there are any limitations in the types of vision tasks it can be effectively applied to.

Overall, the paper presents a promising direction for improving the efficiency of ViT models, and the ShiftAddViT and MoE techniques could be useful contributions to the ongoing efforts in accelerating Transformer-based models and [developing more efficient vision Transformers.

Conclusion

The paper proposes a new model called ShiftAddViT that reparameterizes pre-trained Vision Transformers (ViTs) to replace computationally expensive matrix multiplications with more efficient bitwise shifts and additions. This approach, combined with a novel mixture of experts framework for handling the MLP layers, can achieve significant inference speedups on GPUs without requiring full retraining of the model.

The key contributions of this work are the development of the ShiftAddViT model and the latency-aware load-balancing loss function, which together enable efficient hardware deployment of Transformer-based vision models. These techniques could have important implications for the real-world application of ViTs, making them more practical and accessible for a wide range of use cases.