[论文解读] ShiftAddViT: Mixture of Multiplication Primitives Towards Efficient Vision Transformer

Haoran You, Huihong Shi+2 more|arXiv (Cornell University)|Jun 10, 2023

Advanced Neural Network Applications被引用 6

一句话总结

该论文提出ShiftAddViT，通过使用按位移位和加法的混合乘法原语重新参数化注意力和MLP层，实现无需微调的端到端GPU推理加速。该方法在专家混合（MoE）框架中引入一种延迟感知的负载均衡损失，动态将token路由至高效移位专家或高精度乘法专家，实现最高5.18倍的延迟降低和42.9%的能效节省，同时保持模型精度。

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 $ extbf{ShiftAddViT}$, which aims to achieve end-to-end inference speedups on GPUs without requiring training from scratch. Specifically, all $ exttt{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 $ extbf{5.18$ imes$}$ latency reductions on GPUs and $ extbf{42.9}$% energy savings, while maintaining a comparable accuracy as original or efficient ViTs.

研究动机与目标

• 降低视觉Transformer因注意力和MLP中密集乘法操作导致的高硬件延迟和高能耗。
• 通过重新参数化预训练的ViT，实现在无需从头训练的前提下，实现GPU上的端到端推理加速。
• 通过自适应路由机制将token分配至不同计算路径，在用廉价移位和加法操作替代乘法时保持模型精度。
• 通过利用硬件友好的原语和新颖的负载均衡机制，协同优化注意力和MLP的效率。
• 设计一种专家混合框架，其中专家在计算成本和精度上存在差异，实现基于延迟的动态token分配。

提出的方法

• 在汉明空间中将查询和键映射为二进制编码后，使用加法核重新参数化自注意力中的所有矩阵乘法。
• 通过移位核重新利用剩余的MLP和线性层，将乘法操作替换为按位移位和加法。
• 通过TVM实现自定义CUDA内核，以优化移位和加法操作在GPU上的高效部署。
• 引入一种专家混合（MoE）框架，其中乘法和移位操作作为计算成本不同的非平衡专家。
• 设计一种延迟感知的负载均衡损失，将更多输入token分配给更快的专家（如移位操作），同时为关键token保留精度。
• 利用路由器根据延迟和重要性动态将重要且敏感的token路由至乘法专家，将不关键的token路由至移位专家。

实验结果

研究问题

• RQ1 能否通过用移位和加法原语替代乘法操作，在不微调的前提下加速预训练的视觉Transformer？
• RQ2 在视觉Transformer中，用近似移位-加法操作替代密集乘法时，如何保持模型精度？
• RQ3 具有非平衡专家（乘法 vs. 移位）的专家混合框架，是否能优于固定替换策略，实现更优的精度-效率权衡？
• RQ4 延迟感知的负载均衡损失是否能有效平衡token路由，以最小化同步开销并最大化加速效果？
• RQ5 基于硬件启发的原语（如移位和加法）在2D和3D基于Transformer的视觉模型中，能在多大程度上降低延迟和能耗？

主要发现

• ShiftAddViT在GPU上实现了最高5.18倍的延迟降低，同时保持与原始或高效ViT相当的精度。
• 与基线ViT相比，该方法将能耗降低了42.9%，展现出出色的能效表现。
• 延迟感知的负载均衡损失在PVTv2-B0上实现了14.6%的延迟降低（从100%降至85.4%），仅导致0.01%的精度下降。
• 可视化结果证实，路由器成功将含对象的token分配至乘法专家，将背景token分配至移位专家。
• 通过加法核重新参数化注意力层可保持模型精度，但仅对MLP进行重新参数化会导致精度下降，因此必须引入MoE框架。
• 具有非平衡专家的MoE框架提供了一种软性替换策略，在保持关键token精度的同时，加速非关键token的处理。