ShiftAddViT: Mixture of Multiplication Primitives Towards Efficient Vision Transformer

TL;DR

ShiftAddViT introduces a hardware-inspired reparameterization that replaces dense multiplications in ViTs with a mix of addition primitives and bitwise shifts, starting from pre-trained models to enable end-to-end GPU speedups without full retraining. It reparameterizes attention using Add layers via Q/K binarization and a Q(KV) ordering, and tackles MLP latency via a novel mixture-of-experts framework that assigns expensive multiplications to important tokens and cheap shifts to others, guided by a latency-aware load-balancing loss. The approach is implemented with TVM kernels and evaluated across 2D and 3D vision tasks, delivering up to $5.18\times$ GPU latency reductions and up to $42.9\%$ energy savings on 2D tasks, and substantial gains in 3D NVS with minimal accuracy loss. The work demonstrates that combining hardware-inspired primitives with dynamic expert routing can significantly accelerate ViTs in practice, offering a scalable path toward energy-efficient transformer inference.

Abstract

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 $\textbf{ShiftAddViT}$, which aims to achieve end-to-end inference speedups on GPUs without requiring training from scratch. Specifically, all $\texttt{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 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. Extensive experiments on various 2D/3D Transformer-based vision tasks consistently validate the effectiveness of our proposed ShiftAddViT, achieving up to $\textbf{5.18$\times$}$ latency reductions on GPUs and $\textbf{42.9}$% energy savings, while maintaining a comparable accuracy as original or efficient ViTs.

Figures (10)

• Figure 1: An illustration of the network architecture overview of ShiftAddViT.
• Figure 2: Illustration of both Shift and Add layers, where $\mathbf{X}$ and $\mathbf{O}$ refer to inputs and outputs.
• Figure 3: Energy breakdown on an Eyeriss accelerator.
• Figure 4: MLP/Linear latency speedups using shifts, where inputs are of shape $B \times K \times M$, weights are of shape $K \times N$. All dimensions are set w.r.t. the real dimensions in PVTs PVT_ICCV.
• Figure 5: MatMuls latency speedups using adds, where inputs are of shape $B\times H \times K \times M$, weights are of shape $B \times H \times K \times N$.