ShiftAddViT：面向高效视觉变换器的乘法基元混合

视觉变换器（ViTs）已展现出令人印象深刻的性能，并成为多种视觉任务的统一骨干网络。然而，ViTs 中的注意力机制和多层感知机（MLPs）由于密集的乘法操作而效率不足，导致训练和推理成本高昂。为此，我们提出用乘法基元（如按位移位和加法）的混合对预训练 ViTs 进行重参数化，从而得到一种新型的减少乘法的模型，称为 $\textbf{ShiftAddViT}$，其目标是在 GPU 上实现端到端推理加速，且无需从头训练。具体而言，查询、键和值之间的所有 $\texttt{MatMuls}$ 在将查询和键映射到汉明空间中的二进制码后，使用加性核进行重参数化。其余的 MLPs 或线性层随后用移位核进行重参数化。我们利用 TVM 来实现并优化这些定制核，以在 GPU 上进行实际硬件部署。我们发现，对注意力进行此类重参数化可保持模型精度，而当应用于 MLPs 时不可避免地会导致精度下降。为了兼顾两者优势，我们进一步提出一种新的专家混合（MoE）框架，通过将乘法或其基元作为专家（如乘法和移位）来对 MLPs 进行重参数化，并设计一种新的延迟感知负载均衡损失。该损失有助于训练一个通用路由器，根据延迟将动态数量的输入令牌分配给不同专家。在各种基于 2D/3D Transformer 的视觉任务上的大量实验一致验证了我们所提出的 ShiftAddViT 的有效性，在 GPU 上实现了高达 \textbf{5.18\times} 的延迟降低和 $\textbf{42.9}$% 的能耗节省，同时保持与原始或高效 ViTs 相当的精度。